[jaxlib] Bundle _{tpu,triton,mosaic_gpu}_ext.pyi with jaxlib

PiperOrigin-RevId: 874094949

Author: superbobry

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -3587,16 +3587,18 @@ def tree_unflatten(cls, aux, flat_registers):
 
 @runtime_checkable
 class TransferPlan(Protocol):
-  tile_index_transforms: tuple[IndexTransform, ...]
+  @property
+  def tile_index_transforms(self) -> tuple[IndexTransform, ...]:
+    raise NotImplementedError
 
-  def select(self, group_elems: Sequence[ir.Value]) -> ir.Value:
+  def select(self, group_elems: Sequence[ir.Value], /) -> ir.Value:
     """Selects the value corresponding to the group of the current thread.
 
     The argument must be of the same length as tile_index_transforms.
     """
     raise NotImplementedError
 
-  def select_if_group(self, group_idx: int, old: ir.Value, new: ir.Value) -> ir.Value:
+  def select_if_group(self, group_idx: int, old: ir.Value, new: ir.Value, /) -> ir.Value:
     """Returns `new` if the current thread belongs to the given group and `old` otherwise.
 
     group_idx must be between 0 and len(tile_index_transforms) - 1.

jaxlib/mlir/_mlir_libs/BUILD.bazel

@@ -194,6 +194,7 @@ nanobind_pywrap_extension(
     name = "_mosaic_gpu_ext",
     srcs = ["mosaic_gpu_ext.cc"],
     copts = COPTS,
+    pytype_srcs = ["_mosaic_gpu_ext.pyi"],
     deps = [
         "//jaxlib/mosaic/dialect/gpu:gpu_dialect_capi",
         "//jaxlib/mosaic/gpu:tiled_layout",
@@ -216,6 +217,7 @@ nanobind_pywrap_extension(
     name = "_tpu_ext",
     srcs = ["tpu_ext.cc"],
     copts = COPTS,
+    pytype_srcs = ["_tpu_ext.pyi"],
     deps = [
         "//jaxlib/mosaic:tpu_dialect_capi",
         "@com_google_absl//absl/log:check",

jaxlib/mlir/_mlir_libs/_mosaic_gpu_ext.pyi

@@ -0,0 +1,176 @@
+# Copyright 2025 The JAX Authors.
+#
+# 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
+#
+#     https://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.
+
+from collections.abc import Iterable, Sequence
+import enum
+from mlir import ir
+
+def register_dialect(context: ir.Context, load: bool = ...) -> None: ...
+def register_inliner_extensions(arg: ir.Context, /) -> None: ...
+
+class TileTransformAttr(ir.Attribute):
+  @staticmethod
+  def isinstance(other_attribute: ir.Attribute) -> bool: ...
+  def __repr__(self) -> str: ...
+  @staticmethod
+  def get(
+      tiling: Sequence[int], context: ir.Context | None = None
+  ) -> TileTransformAttr:
+    """Creates a TileTransformAttr with the given tiling."""
+
+  @property
+  def tiling(self) -> list[int]: ...
+
+class TransposeTransformAttr(ir.Attribute):
+  @staticmethod
+  def isinstance(other_attribute: ir.Attribute) -> bool: ...
+  def __repr__(self) -> str: ...
+  @staticmethod
+  def get(
+      permutation: Sequence[int], context: ir.Context | None = None
+  ) -> TransposeTransformAttr:
+    """Creates a TransposeTransformAttr with the given permutation."""
+
+  @property
+  def permutation(self) -> list[int]: ...
+
+class SwizzleTransformAttr(ir.Attribute):
+  @staticmethod
+  def isinstance(other_attribute: ir.Attribute) -> bool: ...
+  def __repr__(self) -> str: ...
+  @staticmethod
+  def get(
+      swizzle: int, context: ir.Context | None = None
+  ) -> SwizzleTransformAttr:
+    """Creates a SwizzleTransformAttr with the given swizzle."""
+
+  @property
+  def swizzle(self) -> int: ...
+
+def init_cc_mlir(arg: object, /) -> bool: ...
+
+class Tiling:
+  def __init__(self, tiles: Iterable) -> None: ...
+  def tile_shape(self, shape: Sequence[int]) -> tuple: ...
+  def untile_shape(self, shape: Sequence[int]) -> tuple: ...
+  def tile_strides(self, strides: Sequence[int]) -> tuple: ...
+  def tile_indices(self, indices: Sequence[int]) -> tuple: ...
+  def untile_indices(self, indices: Sequence[int]) -> tuple: ...
+  def tile_nested_shape_strides(
+      self, shape: Sequence[Sequence[int]], strides: Sequence[Sequence[int]]
+  ) -> tuple: ...
+  def tile_dimension(self, dim: int) -> tuple: ...
+  def remove_dimension(self, dim: int) -> Tiling: ...
+  def canonicalize(self) -> Tiling: ...
+  @property
+  def tiles(self) -> tuple: ...
+  def __str__(self) -> str: ...
+  def __repr__(self) -> str: ...
+  def __eq__(self, other: object) -> bool: ...
+  def __hash__(self) -> int: ...
+
+class Replicated:
+  def __init__(self, times: int) -> None: ...
+  @property
+  def times(self) -> int: ...
+  @times.setter
+  def times(self, arg: int, /) -> None: ...
+  def __repr__(self) -> str: ...
+  def __hash__(self) -> int: ...
+  def __eq__(self, arg: object, /) -> bool: ...
+
+class TiledLayout:
+  def __init__(
+      self,
+      tiling: Tiling,
+      warp_dims: Iterable,
+      lane_dims: Iterable,
+      vector_dim: int,
+      _check_canonical: bool = ...,
+  ) -> None: ...
+  @property
+  def warp_dims(self) -> tuple: ...
+  @property
+  def lane_dims(self) -> tuple: ...
+  @property
+  def partitioned_warp_dims(self) -> tuple: ...
+  @property
+  def partitioned_lane_dims(self) -> tuple: ...
+  @property
+  def vector_length(self) -> int: ...
+  @property
+  def vector_dim(self) -> int: ...
+  @property
+  def tiling(self) -> Tiling: ...
+  @property
+  def tiled_tiling_shape(self) -> tuple: ...
+  @property
+  def tiled_tiling_rank(self) -> int: ...
+  def warp_indices(self) -> tuple: ...
+  def lane_indices(self) -> tuple: ...
+  def canonicalize(self) -> TiledLayout: ...
+  def registers_shape(self, shape: Sequence[int]) -> tuple: ...
+  def registers_element_type(self, t: ir.Type) -> object: ...
+  def shape_from_registers_shape(self, shape: Sequence[int]) -> tuple: ...
+  @property
+  def base_tile_shape(self) -> tuple: ...
+  def remove_dimension(self, dim: int) -> TiledLayout: ...
+  def reduce(self, axes: Iterable) -> TiledLayout: ...
+  def thread_idxs(self, arg: Sequence[int], /) -> list: ...
+  def __str__(self) -> str: ...
+  def __repr__(self) -> str: ...
+  def __hash__(self) -> int: ...
+  def __eq__(self, other: object | None) -> bool: ...
+
+class Rounding(enum.Enum):
+  UP = 0
+
+  DOWN = 1
+
+class TileTransform:
+  def __init__(
+      self, tiling: Sequence[int], rounding: Rounding | None = ...
+  ) -> None: ...
+  def apply(self, arg: object, /) -> ir.Value: ...
+  def transform_index(self, arg: Iterable, /) -> tuple: ...
+  def transform_shape(self, arg: Sequence[int], /) -> tuple: ...
+  def transform_strides(self, arg: Sequence[int], /) -> tuple: ...
+
+class TrivialTransferPlan:
+  def __init__(self) -> None: ...
+  @property
+  def tile_index_transforms(self) -> tuple: ...
+  def select(self, arg: Iterable, /) -> ir.Value: ...
+  def select_if_group(
+      self, arg0: int, arg1: ir.Value, arg2: ir.Value, /
+  ) -> ir.Value: ...
+
+class StaggeredTransferPlan:
+  def __init__(
+      self, stagger: int, dim: int, size: int, group_pred: object
+  ) -> None: ...
+  @property
+  def stagger(self) -> int: ...
+  @property
+  def dim(self) -> int: ...
+  @property
+  def size(self) -> int: ...
+  @property
+  def group_pred(self) -> ir.Value: ...
+  @property
+  def tile_index_transforms(self) -> tuple: ...
+  def select(self, arg: Iterable, /) -> ir.Value: ...
+  def select_if_group(
+      self, arg0: int, arg1: ir.Value, arg2: ir.Value, /
+  ) -> ir.Value: ...

jaxlib/mlir/_mlir_libs/_tpu_ext.pyi

@@ -0,0 +1,34 @@
+# Copyright 2025 The JAX Authors.
+#
+# 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
+#
+#     https://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.
+
+from mlir import ir
+
+def register_dialect(context: ir.Context, load: bool = ...) -> None: ...
+def private_has_communication(
+    arg: ir.Operation, /
+) -> tuple[bool, bool]: ...
+def private_set_arg_attr(
+    arg0: ir.Operation, arg1: int, arg2: str, arg3: ir.Attribute, /
+) -> None: ...
+
+class Float8EXMYType(ir.Type):
+  @staticmethod
+  def isinstance(other_type: ir.Type) -> bool: ...
+  def __repr__(self) -> str: ...
+  @staticmethod
+  def get(
+      exmy_type: ir.Type | None = None, ctx: ir.Context | None = None
+  ) -> Float8EXMYType: ...
+  @property
+  def underlying_type(self) -> ir.Type: ...

jaxlib/mlir/_mlir_libs/_triton_ext.pyi

@@ -1,4 +1,4 @@
-# Copyright 2024 The JAX Authors.
+# Copyright 2025 The JAX Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,19 +17,20 @@ from jaxlib.mlir import ir
 def register_dialect(context: ir.Context, load: bool = ...) -> None: ...
 
 class PointerType(ir.Type):
-  @classmethod
-  def get(cls, pointee_type: ir.Type, address_space: int) -> PointerType: ...
-
   @staticmethod
-  def isinstance(other: ir.Type) -> bool: ...
+  def isinstance(other_type: ir.Type) -> bool: ...
+  def __repr__(self) -> str: ...
+  @staticmethod
+  def get_static_typeid() -> ir.TypeID: ...
+  @staticmethod
+  def get(pointee_type: ir.Type, address_space: int) -> PointerType:
+    """Creates a PointerType type."""
 
   @property
   def pointee_type(self) -> ir.Type: ...
-
   @property
   def address_space(self) -> int: ...
 
 def infer_reduce_op_encoding(
-    op_attribute: ir.Attribute,
-    axis: int,
-) -> ir.Attribute: ...
+    arg0: ir.Attribute, arg1: int, /
+) -> ir.Attribute | None: ...