Merge pull request #293 from superbobry:main

PiperOrigin-RevId: 675154784

.github/workflows/ci.yml

@@ -3,15 +3,17 @@ on:
   push:
     branches:
       - main
-permissions:
-  contents: write
+  pull_request:
+    branches:
+      - main
+
 jobs:
-  deploy:
+  lint:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # ratchet:actions/checkout@v4
+      - name: Set up Python 3.10
+        uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # ratchet:actions/setup-python@v5
         with:
-          python-version: 3.x
-      - run: pip install -r docs/requirements.txt
-      - run: mkdocs gh-deploy --force
+          python-version: '3.10'
+      - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # ratchet:pre-commit/action@v3.0.1

.github/workflows/docs.yml

@@ -0,0 +1,20 @@
+name: docs
+on:
+  push:
+    branches:
+      - main
+
+permissions:
+  contents: write
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # ratchet:actions/checkout@v4
+      - name: Set up Python 3.10
+        uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # ratchet:actions/setup-python@v5
+        with:
+          python-version: '3.10'
+      - run: pip install -r docs/requirements.txt
+      - run: mkdocs gh-deploy --force

.pre-commit-config.yaml

@@ -0,0 +1,31 @@
+# Install the pre-commit hooks below with
+# 'pre-commit install'
+
+# Auto-update the version of the hooks with
+# 'pre-commit autoupdate'
+
+# Run the hooks on all files with
+# 'pre-commit run --all'
+
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: 2c9f875913ee60ca25ce70243dc24d5b6415598c  # frozen: v4.6.0
+  hooks:
+  - id: check-ast
+  - id: check-merge-conflict
+  - id: check-toml
+  - id: check-yaml
+  - id: end-of-file-fixer
+    # only include python files
+    files: \.py$
+  - id: debug-statements
+    # only include python files
+    files: \.py$
+  - id: trailing-whitespace
+    # only include python files
+    files: \.py$
+
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: 8b5112a3b2ad121439a2092f8ff548c0d80f2514  # frozen: v0.6.1
+  hooks:
+  - id: ruff

examples/block_map.py

@@ -18,7 +18,6 @@
 
 import jax
 import jax.numpy as jnp
-from jax import lax
 from jax import random
 import jax_triton as jt
 from jax_triton import pallas as pl
@@ -178,8 +177,8 @@ def main(unused_argv):
   k = random.normal(k_key, shape, dtype=dtype)
   v = random.normal(v_key, shape, dtype=dtype)
 
-  o = mha(q, k, v).block_until_ready()
-  o_ref = mha_reference(q, k, v).block_until_ready()
+  mha(q, k, v).block_until_ready()
+  mha_reference(q, k, v).block_until_ready()
 
 if __name__ == "__main__":
   from absl import app

examples/pallas/blocksparse_matmul.py

@@ -21,7 +21,6 @@
 from jax import random
 import jax
 from jax import lax
-import jax.numpy as jnp
 import numpy as np
 
 import jax_triton as jt
@@ -87,7 +86,7 @@ def tree_unflatten(cls, data, xs):
     return BlockELL(blocks, blocks_per_row, indices, shape=shape)
 
   def _validate(self):
-    nblocks, n, m = self.blocks.shape
+    _nblocks, n, m = self.blocks.shape
     nrows = self.blocks_per_row.shape[0]
     assert self.indices.shape[0] == nrows
     assert len(self.shape) == 2
@@ -168,7 +167,7 @@ def sdd_matmul(x_ell, y, num_warps: int = 8, num_stages: int = 3, bn: int = 64,
   grid = (jt.cdiv(m, bm), jt.cdiv(n, bn))
 
   kernel = functools.partial(sdd_kernel, bm=bm, bn=bn)
-  out_shape = jax.ShapeDtypeStruct(shape=(m, n), dtype=x.dtype)
+  out_shape = jax.ShapeDtypeStruct(shape=(m, n), dtype=x_ell.dtype)
   return pl.pallas_call(kernel, num_warps=num_warps, num_stages=num_stages,
                         grid=grid, out_shape=out_shape,
                         debug=debug)(x_ell.blocks, x_ell.indices,

examples/pallas/lstm.py

@@ -17,14 +17,11 @@
 import functools
 import timeit
 
-from typing import Optional, Tuple
 
 import jax.numpy as jnp
 from jax import random
 import jax
-from jax import lax
 from jax._src.lax.control_flow import for_loop
-import jax.numpy as jnp
 import numpy as np
 
 import jax_triton as jt
@@ -188,13 +185,15 @@ def main(unused_argv):
   x = random.normal(x_key, (batch_size, feature_size), dtype)
   h = random.normal(h_key, (batch_size, hidden_size), dtype)
   c = random.normal(c_key, (batch_size, hidden_size), dtype)
-  lstm_cell = jax.jit(functools.partial(lstm_cell,
-                                        block_batch=block_batch,
-                                        block_hidden=block_hidden,
-                                        block_features=block_features,
-                                        num_warps=num_warps,
-                                        num_stages=num_stages))
-  y, c_next = jax.block_until_ready(lstm_cell(weights, x, h, c))
+  lstm_cell_fn = jax.jit(functools.partial(
+      lstm_cell,
+      block_batch=block_batch,
+      block_hidden=block_hidden,
+      block_features=block_features,
+      num_warps=num_warps,
+      num_stages=num_stages,
+  ))
+  y, c_next = jax.block_until_ready(lstm_cell_fn(weights, x, h, c))
   y_ref, c_next_ref = lstm_cell_reference(weights, x, h, c)
   np.testing.assert_allclose(y, y_ref, atol=0.05, rtol=0.05)
   np.testing.assert_allclose(c_next, c_next_ref, atol=0.05, rtol=0.05)

jax_triton/__init__.py

@@ -13,6 +13,17 @@
 # limitations under the License.
 
 """Library for JAX-Triton integrations."""
+
+__all__ = [
+    "utils",
+    "triton_call",
+    "cdiv",
+    "next_power_of_2",
+    "strides_from_shape",
+    "__version__",
+    "__version_info__",
+]
+
 import jaxlib
 from jax._src.lib import gpu_triton
 from jax_triton import utils

jax_triton/experimental/fusion/__init__.py

@@ -19,4 +19,4 @@
 jax.nn.sigmoid = sigmoid
 del sigmoid, oryx, jax
 
-from jax_triton.experimental.fusion.lowering import jit
+from jax_triton.experimental.fusion.lowering import jit as jit

jax_triton/experimental/fusion/fusion.py

@@ -23,7 +23,6 @@
 from jax import lax
 from jax.extend import linear_util as lu
 from jax.interpreters import partial_eval as pe
-from jax.interpreters import xla
 from jax._src import core
 from jax._src import util
 from jax._src.lax.control_flow import for_loop
@@ -251,6 +250,7 @@ def _matmul_elementwise_lowering_rule(x, y, *args, left_ops, right_ops, out_ops,
     bias, = args
   else:
     bias = None
+  del bias  # TODO(sharadmv): Please fix or remove `bias` above.
   lhs_dim, rhs_dim = contract_dims
   M, N, K = x.shape[1 - lhs_dim], y.shape[1 - rhs_dim], x.shape[lhs_dim]
   assert x.shape[lhs_dim] == y.shape[rhs_dim]
@@ -340,4 +340,3 @@ def _dot_general_lowering_rule(x, y, dimension_numbers, **_):
                                            out_ops=[], contract_dims=(lhs_dim,
                                                                       rhs_dim))
 lowering_rules[lax.dot_general_p] = _dot_general_lowering_rule
-

jax_triton/experimental/fusion/jaxpr_rewriter.py

@@ -19,7 +19,7 @@
 import dataclasses
 import itertools as it
 
-from typing import Any, Callable, Dict, List, Set, Tuple, Union
+from typing import Any, Callable, List, Tuple, Union
 
 from jax._src import core as jax_core
 import jax.numpy as jnp

jax_triton/triton_lib.py

@@ -336,8 +336,6 @@ def compile_ttir_to_hsaco_inplace(
   amdgcn = hip_backend.make_amdgcn(llir, metadata, hip_options)
   hsaco = hip_backend.make_hsaco(amdgcn, metadata, hip_options)
 
-  if hip_options.debug:
-    print(x)
   name = metadata["name"]
   ttgir = str(ttgir) if _JAX_TRITON_DUMP_DIR else None
   llir = str(llir) if _JAX_TRITON_DUMP_DIR else None

jax_triton/utils.py

@@ -13,6 +13,11 @@
 # limitations under the License.
 
 """Contains utilities for writing and calling Triton functions."""
+
+
+__all__ = ["cdiv", "strides_from_shape", "next_power_of_2"]
+
+
 from jax.experimental.pallas import cdiv
 from jax.experimental.pallas import strides_from_shape
-from jax.experimental.pallas import next_power_of_2
+from jax.experimental.pallas import next_power_of_2

pyproject.toml

@@ -26,3 +26,25 @@ packages = ["jax_triton"]
 
 [tool.setuptools.dynamic]
 version = {attr = "jax_triton.version.__version__"}
+
+[tool.ruff]
+preview = true
+exclude = [
+    ".git",
+    "build",
+    "__pycache__",
+    "*.ipynb",
+]
+line-length = 88
+indent-width = 2
+target-version = "py310"
+
+[tool.ruff.lint]
+ignore = [
+    # Do not assign a `lambda` expression, use a `def`
+    "E731",
+    # Module level import not at top of file
+    "E402",
+    # Ambiguous variable name
+    "E741",
+]