Merge pull request #295 from superbobry:main

PiperOrigin-RevId: 675199892

examples/block_map.py

@@ -14,7 +14,6 @@
 
 import functools
 
-from typing import Optional
 
 import jax
 import jax.numpy as jnp
@@ -113,7 +112,7 @@ def mha(q, k, v, *,
     sm_scale: float = 1.0,
     block_q: int = 128,
     block_k: int = 128,
-    num_warps: Optional[int] = None,
+    num_warps: int | None = None,
     num_stages: int = 1,
     grid=None,
     ):

examples/pallas/blocksparse_matmul.py

@@ -15,7 +15,6 @@
 import functools
 import timeit
 
-from typing import Tuple
 
 import jax.numpy as jnp
 from jax import random
@@ -63,7 +62,7 @@ class BlockELL:
   blocks: jnp.ndarray  # float32[n_rows, n_blocks, *block_size]
   blocks_per_row: jnp.ndarray  # int32[n_rows, n_blocks]
   indices: jnp.ndarray  # int32[n_rows, max_num_blocks_per_row, 2]
-  shape: Tuple[int, int] # (n_rows * block_size[0], n_cols * block_size[1])
+  shape: tuple[int, int] # (n_rows * block_size[0], n_cols * block_size[1])
 
   ndim: int = property(lambda self: len(self.shape))
   num_blocks = property(lambda self: self.blocks.shape[0])

examples/pallas/lstm.py

@@ -67,7 +67,7 @@ def body(k, acc_refs):
   accs = for_loop.for_loop(num_k_blocks, body, [acc_i, acc_f, acc_o, acc_g])
   bs = [pl.load(b_ref, (idx_n,))
         for b_ref in [b_hi_ref, b_hf_ref, b_hg_ref, b_ho_ref]]
-  acc_i, acc_f, acc_g, acc_o = [acc + b for acc, b in zip(accs, bs)]
+  acc_i, acc_f, acc_g, acc_o = (acc + b for acc, b in zip(accs, bs))
   i_gate, f_gate, o_gate = (
       jax.nn.sigmoid(acc_i), jax.nn.sigmoid(acc_f), jax.nn.sigmoid(acc_o))
   cell = jnp.tanh(acc_g)
@@ -124,7 +124,7 @@ def lstm_cell_reference(weights, x, h, c):
   xs = [jnp.dot(x, w) for w in ws]
   hs = [jnp.dot(h, u) for u in us]
   accs = [x + h for x, h in zip(xs, hs)]
-  acc_i, acc_f, acc_g, acc_o = [acc + b[None] for acc, b in zip(accs, bs)]
+  acc_i, acc_f, acc_g, acc_o = (acc + b[None] for acc, b in zip(accs, bs))
   i_gate, f_gate, o_gate = (
       jax.nn.sigmoid(acc_i), jax.nn.sigmoid(acc_f), jax.nn.sigmoid(acc_o))
   cell = jnp.tanh(acc_g)

jax_triton/experimental/fusion/fusion.py

@@ -17,7 +17,7 @@
 import functools
 import os
 
-from typing import Any, Tuple
+from typing import Any
 
 import jax
 from jax import lax
@@ -204,7 +204,7 @@ def make_elementwise(shape, dtype, *args):
 class MatmulElementwise(jax_rewrite.JaxExpression):
   x: jax_rewrite.JaxExpression
   y: jax_rewrite.JaxExpression
-  elem_ops: Tuple[core.Primitive]
+  elem_ops: tuple[core.Primitive]
 
   def match(self, expr, bindings, succeed):
     if not isinstance(expr, MatmulElementwise):

jax_triton/experimental/fusion/jaxpr_rewriter.py

@@ -19,7 +19,8 @@
 import dataclasses
 import itertools as it
 
-from typing import Any, Callable, List, Tuple, Union
+from typing import Any
+from collections.abc import Callable
 
 from jax._src import core as jax_core
 import jax.numpy as jnp
@@ -35,7 +36,7 @@
 class Node(matcher.Pattern, metaclass=abc.ABCMeta):
 
   @abc.abstractproperty
-  def parents(self) -> List[Node]:
+  def parents(self) -> list[Node]:
     ...
 
 
@@ -51,9 +52,9 @@ def map_parents(self, fn: Callable[[Node], Node]) -> Node:
 class Eqn(Node):
   primitive: jax_core.Primitive
   params: jr.Params
-  invars: List[Node]
-  shape: Union[Tuple[int, ...], List[Tuple[int, ...]]]
-  dtype: Union[jnp.dtype, List[jnp.dtype]]
+  invars: list[Node]
+  shape: tuple[int, ...] | list[tuple[int, ...]]
+  dtype: jnp.dtype | list[jnp.dtype]
 
   @property
   def parents(self):
@@ -77,7 +78,7 @@ def match(self, expr, bindings, succeed):
 
 @dataclasses.dataclass(frozen=True, eq=False)
 class JaxprVar(Node):
-  shape: Tuple[int, ...]
+  shape: tuple[int, ...]
   dtype: jnp.dtype
 
   def match(self, expr, bindings, succeed):
@@ -131,7 +132,7 @@ def from_literal(cls, var: jax_core.Literal) -> Literal:
 @dataclasses.dataclass(eq=False)
 class Part(Node):
   index: int
-  shape: Tuple[int, ...]
+  shape: tuple[int, ...]
   dtype: jnp.dtype
   parent: Node
 
@@ -153,9 +154,9 @@ def map_parents(self, fn):
 
 @dataclasses.dataclass(eq=True)
 class JaxprGraph(matcher.Pattern):
-  constvars: List[Node]
-  invars: List[Node]
-  outvars: List[Node]
+  constvars: list[Node]
+  invars: list[Node]
+  outvars: list[Node]
 
   def get_nodes(self):
     nodes = set(self.outvars)
@@ -167,7 +168,7 @@ def get_nodes(self):
         queue.append(p)
     return nodes
 
-  def get_children(self, node) -> List[Node]:
+  def get_children(self, node) -> list[Node]:
     nodes = self.get_nodes()
     return [n for n in nodes if node in n.parents]
 
@@ -274,7 +275,7 @@ def to_jaxpr(self) -> jax_core.Jaxpr:
     outvars = [env[n] for n in self.outvars]
     return jax_core.Jaxpr(constvars, invars, outvars, eqns, jax_core.no_effects)
 
-  def toposort(self) -> List[Node]:
+  def toposort(self) -> list[Node]:
     node_stack = list(self.outvars)
     child_counts = {}
     while node_stack:

jax_triton/experimental/fusion/lowering.py

@@ -15,7 +15,7 @@
 """Contains lowering passes for jaxprs to pallas."""
 import functools
 
-from typing import Any, Dict
+from typing import Any
 
 import jax
 from jax import api_util
@@ -317,7 +317,7 @@ def read(v: core.Atom) -> Any:
   def write(v: Var, val: Any) -> None:
     env[v] = val
 
-  env: Dict[Var, Any] = {}
+  env: dict[Var, Any] = {}
   map(write, jaxpr.constvars, consts)
   map(write, jaxpr.invars, args)
   for eqn in jaxpr.eqns:

jax_triton/triton_lib.py

@@ -25,7 +25,8 @@
 import pprint
 import tempfile
 import types
-from typing import Any, Callable, Dict, Optional, Protocol, Sequence, Tuple, Union
+from typing import Any, Protocol, Union
+from collections.abc import Callable, Sequence
 import zlib
 from functools import partial
 
@@ -102,11 +103,11 @@
     jnp.dtype("bool"): "B",
 }
 
-Grid = Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]]
-GridOrLambda = Union[Grid, Callable[[Dict[str, Any]], Grid]]
+Grid = Union[int, tuple[int], tuple[int, int], tuple[int, int, int]]
+GridOrLambda = Union[Grid, Callable[[dict[str, Any]], Grid]]
 
 
-def normalize_grid(grid: GridOrLambda, metaparams) -> Tuple[int, int, int]:
+def normalize_grid(grid: GridOrLambda, metaparams) -> tuple[int, int, int]:
   if callable(grid):
     grid = grid(metaparams)
   if isinstance(grid, int):
@@ -186,8 +187,8 @@ class CompilationResult:
   name: str
   shared_mem_bytes: int
   cluster_dims: tuple
-  ttgir: Optional[str]
-  llir: Optional[str]
+  ttgir: str | None
+  llir: str | None
 
 def compile_ttir_inplace(
     ttir,
@@ -375,7 +376,7 @@ def get_or_create_triton_kernel(
     enable_fp_fusion,
     metaparams,
     dump: bool,
-) -> Tuple[triton_kernel_call_lib.TritonKernel, Any]:
+) -> tuple[triton_kernel_call_lib.TritonKernel, Any]:
   if num_warps is None:
     num_warps = 4
   if num_stages is None:
@@ -730,7 +731,7 @@ def prune_configs(configs, named_args, **kwargs):
 class ShapeDtype(Protocol):
 
   @property
-  def shape(self) -> Tuple[int, ...]:
+  def shape(self) -> tuple[int, ...]:
     ...
 
   @property
@@ -739,21 +740,21 @@ def dtype(self) -> np.dtype:
 
 
 def triton_call(
-    *args: Union[jax.Array, bool, int, float, np.float32],
+    *args: jax.Array | bool | int | float | np.float32,
     kernel: triton.JITFunction,
-    out_shape: Union[ShapeDtype, Sequence[ShapeDtype]],
+    out_shape: ShapeDtype | Sequence[ShapeDtype],
     grid: GridOrLambda,
     name: str = "",
     custom_call_target_name: str = "triton_kernel_call",
-    num_warps: Optional[int] = None,
-    num_stages: Optional[int] = None,
+    num_warps: int | None = None,
+    num_stages: int | None = None,
     num_ctas: int = 1,  # TODO(giorgioa): Add support for dimensions tuple.
-    compute_capability: Optional[int] = None,
+    compute_capability: int | None = None,
     enable_fp_fusion: bool = True,
-    input_output_aliases: Optional[Dict[int, int]] = None,
-    zeroed_outputs: Union[
-        Sequence[int], Callable[[Dict[str, Any]], Sequence[int]]
-    ] = (),
+    input_output_aliases: dict[int, int] | None = None,
+    zeroed_outputs: (
+        Sequence[int] | Callable[[dict[str, Any]], Sequence[int]]
+    ) = (),
     debug: bool = False,
     serialized_metadata: bytes = b"",
     **metaparams: Any,

pyproject.toml

@@ -8,7 +8,6 @@ dependencies = [
   "absl-py>=1.4.0",
   "jax>=0.4.31",
   "triton>=3.0",
-  "setuptools",  # triton seems to need this when installing itself.
 ]
 
 [project.optional-dependencies]