Migrate the exteded StableHLO dialect

pyproject.toml

@@ -65,6 +65,14 @@ ignore = [
     "PYI041", # https://docs.astral.sh/ruff/rules/redundant-numeric-union
 ]
 
+[tool.ruff.lint.per-file-ignores]
+# Allow long lines in StableHLO dialect files (contain MLIR examples in docstrings)
+"src/xdsl_jax/dialects/stablehlo/**/*.py" = ["E501"] 
+# Allow long lines in xdsl_extras files (contain detailed docstrings)
+"src/xdsl_jax/xdsl_extras/**/*.py" = ["E501"]
+# Allow long lines in test files (contain MLIR test cases)
+"tests/**/*.py" = ["E501"]
+
 [tool.ruff.lint.mccabe]
 max-complexity = 10
 

src/xdsl_jax/dialects/stablehlo/__init__.py

@@ -0,0 +1,158 @@
+# Copyright 2025 Xanadu Quantum Technologies Inc.
+#
+# 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.
+
+"""
+StableHLO dialect package for xdsl-jax.
+
+This package contains organized elementwise operations and other StableHLO-related
+functionality.
+"""
+
+# Import all elementwise operations explicitly
+from .attributes import (
+    CustomCallApiVersion,
+    CustomCallApiVersionAttr,
+    GatherDimensionNumbers,
+    OutputOperandAlias,
+    ResultAccuracyModeAttr,
+    ScatterDimensionNumbers,
+)
+from .control_flow import (
+    IfOp,
+    OptimizationBarrierOp,
+    WhileOp,
+)
+from .data_movement import (
+    BroadcastInDimOp,
+    ConcatenateOp,
+    DynamicSliceOp,
+    GatherOp,
+    ReshapeOp,
+    ScatterOp,
+    SliceOp,
+)
+
+# Import the main StableHLO dialect
+from .dialect import StableHLO
+from .dynamism import (
+    DynamicBroadcastInDimOp,
+)
+from .elementwise_binary import (
+    ComplexOp,
+    DivideOp,
+    MaximumOp,
+    MinimumOp,
+    PowerOp,
+    RemainderOp,
+)
+from .elementwise_other import (
+    ClampOp,
+    CompareOp,
+    MapOp,
+    ReducePrecisionOp,
+    SelectOp,
+)
+from .elementwise_unary import (
+    ConvertOp,
+    CosineOp,
+    ExponentialMinusOneOp,
+    ExponentialOp,
+    FloorOp,
+    ImagOp,
+    IsFiniteOp,
+    LogisticOp,
+    LogOp,
+    LogPlusOneOp,
+    NegateOp,
+    RealOp,
+    RoundNearestAfzOp,
+    RoundNearestEvenOp,
+    RsqrtOp,
+    SignOp,
+    SineOp,
+    SqrtOp,
+    TanhOp,
+    TanOp,
+)
+from .extensibility import (
+    CustomCallOp,
+)
+from .reduction import (
+    ReduceOp,
+)
+
+# Export all operations and the dialect for external use
+__all__ = [
+    # Main dialect
+    "StableHLO",
+    # Elementwise unary operations
+    "ConvertOp",
+    "CosineOp",
+    "ExponentialMinusOneOp",
+    "ExponentialOp",
+    "FloorOp",
+    "ImagOp",
+    "IsFiniteOp",
+    "LogOp",
+    "LogPlusOneOp",
+    "LogisticOp",
+    "NegateOp",
+    "RealOp",
+    "RoundNearestAfzOp",
+    "RoundNearestEvenOp",
+    "RsqrtOp",
+    "SignOp",
+    "SineOp",
+    "SqrtOp",
+    "TanOp",
+    "TanhOp",
+    # Elementwise binary operations
+    "ComplexOp",
+    "DivideOp",
+    "MaximumOp",
+    "MinimumOp",
+    "PowerOp",
+    "RemainderOp",
+    # Elementwise other operations
+    "ClampOp",
+    "CompareOp",
+    "MapOp",
+    "ReducePrecisionOp",
+    "SelectOp",
+    # Control flow operations
+    "IfOp",
+    "WhileOp",
+    "OptimizationBarrierOp",
+    # Data movement operations
+    "BroadcastInDimOp",
+    "ConcatenateOp",
+    "DynamicSliceOp",
+    "GatherOp",
+    "ReshapeOp",
+    "ScatterOp",
+    "SliceOp",
+    # Dynamism operations
+    "DynamicBroadcastInDimOp",
+    # Reduction operations
+    "ReduceOp",
+    # Extensibility operations
+    "CustomCallOp",
+    # Attributes
+    "GatherDimensionNumbers",
+    "ResultAccuracyModeAttr",
+    "ScatterDimensionNumbers",
+    "CustomCallApiVersion",
+    "CustomCallApiVersionAttr",
+    "OutputOperandAlias",
+]