Add Clang-format checker GitHub Action

.github/workflows/check-clang-format.yaml

@@ -0,0 +1,29 @@
+# A workflow to test clang-format composite action
+
+name: Test CI Clang-format Action
+# Run on pull_request
+on:
+  pull_request:
+    paths:
+    - ci_clangformat/**
+    - tests/ci_clangformat/**
+    - .github/workflows/check-clang-format.yaml
+    branches:
+      - main
+defaults:
+  run:
+    shell: bash
+jobs:
+  run-clang-format:
+    runs-on: "linux-x86-n2-16"
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
+    timeout-minutes: 10
+    if: |
+      github.event.sender.type == 'User' ||
+      contains(github.event.pull_request.body, 'RUN_CLANG_FORMAT')
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          persist-credentials: false
+      - name: "Run clang-format action"
+        uses: ./ci_clangformat/

ci_clangformat/.clang-format.default

@@ -0,0 +1,6 @@
+BasedOnStyle: Google
+Language: Cpp
+PointerBindsToType: true
+SortIncludes: Never
+AlignTrailingComments:
+  Kind: Always

ci_clangformat/README.md

@@ -0,0 +1,27 @@
+# CI Clang-format
+
+This composite action helps maintain consistent C/C++ code style by running
+`clang-format` on modified files in your pull requests. It checks for
+formatting violations and will cause the workflow to fail if any issues are
+found, ensuring code quality before merging.
+
+The action uses your .clang-format style file if present in the repository
+root; otherwise, it will use the .clang-format.default under this folder.
+
+This action offers the following configuration through its inputs:
+* `clang_format_version`: Choose the exact clang-format version to use,
+ with `20.1.5` as the default to align with recent stable releases.
+
+## Resolving Formatting Failures
+If a workflow run fails due to formatting violations, you're expected to
+fix the issues locally. Simply run `clang-format` on the problematic
+files, e.g., using
+`uvx clang-format==20.1.5 -i --verbose --style=file <files>`,
+and then commit the formatted code to your pull request.
+
+## UV Requirement
+This action leverages `uv` to reliably install and run specific
+`clang-format` versions, ensuring consistent behavior across different
+environments. `uvx` is a convenience alias that calls `uv tool run`.
+If `uv` does not exist, you'll need to include a step to [install](https://docs.astral.sh/uv/getting-started/installation/)
+it in your workflow's running environment.

ci_clangformat/action.yaml

@@ -0,0 +1,54 @@
+# Copyright 2025 Google LLC
+
+# 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.
+name: "Check clang-format"
+description: 'Action to run clang-format on changed files'
+inputs:
+  clang_format_version:
+    description: 'The clang-format version to use.'
+    required: true
+    default: "20.1.5"
+
+runs:
+  using: "composite"
+  steps:
+  - name: "Checking out repository"
+    uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+    with:
+      fetch-depth: 0 # Fetch full history for accurate diffing
+  - name: Run clang-format check
+    id: check-format
+    shell: bash
+    env:
+      REPO_CLANG_FORMAT: .clang-format
+      ACTION_DEFAULT_CLANG_FORMAT: ${{ github.action_path }}.clang-format.default
+      CLANG_FORMAT_VERSION: ${{ inputs.clang_format_version }}
+      TARGET_BRANCH: ${{ github.base_ref }}
+    run: |
+      FILE_PATTERNS="*.cc *.h"
+      CLANG_FORMAT_COMMON_ARGS="--dry-run --Werror --verbose"
+
+      # Add this line to resolve the dubious ownership error
+      git config --global --add safe.directory "$GITHUB_WORKSPACE"
+
+      # Compare PR head against the base branch to find changed files.
+      GIT_DIFF_CMD="git diff -z --name-only --diff-filter=d origin/$TARGET_BRANCH HEAD -- $FILE_PATTERNS"
+
+      if [ -f "$REPO_CLANG_FORMAT" ]; then
+        echo "::notice::Using repository's .clang-format file."
+        # uvx (an alias for `uv tool run`) installs and runs clang-format with a specific version.
+        $GIT_DIFF_CMD | xargs -0 uvx clang-format==$CLANG_FORMAT_VERSION $CLANG_FORMAT_COMMON_ARGS
+      else
+        echo "::notice::Repository does not have a .clang-format file. Using the action's default under $ACTION_DEFAULT_CLANG_FORMAT."
+        $GIT_DIFF_CMD | xargs -0 uvx clang-format==$CLANG_FORMAT_VERSION -style=file:$ACTION_DEFAULT_CLANG_FORMAT $CLANG_FORMAT_COMMON_ARGS
+      fi

tests/ci_clangformat/absl_status_casters.h

@@ -0,0 +1,213 @@
+/* Copyright 2023 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
+
+    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.
+==============================================================================*/
+
+#ifndef JAXLIB_ABSL_STATUS_CASTERS_H_
+#define JAXLIB_ABSL_STATUS_CASTERS_H_
+
+#include <stdexcept>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+
+namespace jax {
+
+// C++ -> Python caster helpers.
+//
+// Failing statuses become Python exceptions; OK Status() becomes None.
+//
+// For example:
+//
+// - Functions without arguments:
+//   m.def("my_func", []() { ThrowIfError(MyFunc()); }
+// - Classes with a single argument:
+//   py_class.def("delete", [](Buffer& self) {
+//     ThrowIfError(self.Delete());
+//   }
+//
+// For functions with more arguments, you can either inline the arguments,
+// or use the `ThrowIfErrorWrapper` wrapper defined below:
+//
+// m.def("my_func", ThrowIfErrorWrapper(MyFunc));
+//
+// Nonstatic member functions can be wrapped by passing a
+// pointer-to-member-function:
+// ThrowIfErrorWrapper(&MyClass::MyMethod)
+
+inline void ThrowIfError(absl::Status src) {
+  if (!src.ok()) {
+    throw std::runtime_error(src.ToString());
+  }
+}
+
+// If one does not want to have to define a lambda specifying the inputs
+// arguments, on can use the `ThrowIfErrorWrapper` wrapper.
+//
+// There are three specializations:
+// - For free functions, `Sig` is the function type and `F` is `Sig&`.
+// - For callable types, `Sig` is the pointer to member function type
+//   and `F` is the type of the callable.
+// - For a nonstatic member function of a class `C`, `Sig` is the function type
+//   and `F` is Sig C::*.
+//
+// In the first two cases, the wrapper returns a callable with signature `Sig`;
+// in the third case, the wrapper returns callable with a modified signature
+// that takes a C instance as the first argument.
+template <typename Sig, typename F>
+struct ThrowIfErrorWrapper;
+
+// C++17 "deduction guide" that guides class template argument deduction (CTAD)
+// For free functions.
+template <typename F>
+ThrowIfErrorWrapper(F) -> ThrowIfErrorWrapper<decltype(&F::operator()), F>;
+
+// For callable types (with operator()).
+template <typename... Args>
+ThrowIfErrorWrapper(absl::Status (&)(Args...))
+    -> ThrowIfErrorWrapper<absl::Status(Args...), absl::Status (&)(Args...)>;
+
+// For unbound nonstatic member functions.
+template <typename C, typename... Args>
+ThrowIfErrorWrapper(absl::Status (C::*)(Args...))
+    -> ThrowIfErrorWrapper<absl::Status(Args...), C>;
+
+// Template specializations.
+
+// For free functions.
+template <typename... Args>
+struct ThrowIfErrorWrapper<absl::Status(Args...), absl::Status (&)(Args...)> {
+  explicit ThrowIfErrorWrapper(absl::Status (&f)(Args...)) : func(f) {}
+  void operator()(Args... args) {
+    ThrowIfError(func(std::forward<Args>(args)...));
+  }
+  absl::Status (&func)(Args...);
+};
+
+// For callable types (with operator()), non-const and const versions.
+template <typename C, typename... Args, typename F>
+struct ThrowIfErrorWrapper<absl::Status (C::*)(Args...), F> {
+  explicit ThrowIfErrorWrapper(F &&f) : func(std::move(f)) {}
+  void operator()(Args... args) {
+    ThrowIfError(func(std::forward<Args>(args)...));
+  }
+  F func;
+};
+template <typename C, typename... Args, typename F>
+struct ThrowIfErrorWrapper<absl::Status (C::*)(Args...) const, F> {
+  explicit ThrowIfErrorWrapper(F &&f) : func(std::move(f)) {}
+  void operator()(Args... args) const {
+    ThrowIfError(func(std::forward<Args>(args)...));
+  }
+  F func;
+};
+
+// For unbound nonstatic member functions, non-const and const versions.
+// `ptmf` stands for "pointer to member function".
+template <typename C, typename... Args>
+struct ThrowIfErrorWrapper<absl::Status(Args...), C> {
+  explicit ThrowIfErrorWrapper(absl::Status (C::*ptmf)(Args...)) : ptmf(ptmf) {}
+  void operator()(C &instance, Args... args) {
+    ThrowIfError((instance.*ptmf)(std::forward<Args>(args)...));
+  }
+  absl::Status (C::*ptmf)(Args...);
+};
+template <typename C, typename... Args>
+struct ThrowIfErrorWrapper<absl::Status(Args...) const, C> {
+  explicit ThrowIfErrorWrapper(absl::Status (C::*ptmf)(Args...) const)
+      : ptmf(ptmf) {}
+  void operator()(const C &instance, Args... args) const {
+    ThrowIfError((instance.*ptmf)(std::forward<Args>(args)...));
+  }
+  absl::Status (C::*ptmf)(Args...) const;
+};
+
+// Utilities for `StatusOr`.
+template <typename T>
+T ValueOrThrow(absl::StatusOr<T> v) {
+  if (!v.ok()) {
+    throw std::runtime_error(v.status().ToString());
+  }
+  return std::move(v).value();
+}
+
+template <typename Sig, typename F>
+struct ValueOrThrowWrapper;
+
+template <typename F>
+ValueOrThrowWrapper(F) -> ValueOrThrowWrapper<decltype(&F::operator()), F>;
+
+template <typename R, typename... Args>
+ValueOrThrowWrapper(absl::StatusOr<R> (&)(Args...))
+    -> ValueOrThrowWrapper<absl::StatusOr<R>(Args...),
+                           absl::StatusOr<R> (&)(Args...)>;
+
+template <typename C, typename R, typename... Args>
+ValueOrThrowWrapper(absl::StatusOr<R> (C::*)(Args...))
+    -> ValueOrThrowWrapper<absl::StatusOr<R>(Args...), C>;
+
+// Deduction guide for const methods.
+template <typename C, typename R, typename... Args>
+ValueOrThrowWrapper(absl::StatusOr<R> (C::*)(Args...) const)
+    -> ValueOrThrowWrapper<absl::StatusOr<R>(Args...) const, C>;
+
+template <typename R, typename... Args>
+struct ValueOrThrowWrapper<absl::StatusOr<R>(Args...),
+                           absl::StatusOr<R> (&)(Args...)> {
+  explicit ValueOrThrowWrapper(absl::StatusOr<R> (&f)(Args...)) : func(f) {}
+  R operator()(Args... args) const {
+    return ValueOrThrow(func(std::forward<Args>(args)...));
+  }
+  absl::StatusOr<R> (&func)(Args...);
+};
+template <typename R, typename C, typename... Args, typename F>
+struct ValueOrThrowWrapper<absl::StatusOr<R> (C::*)(Args...), F> {
+  explicit ValueOrThrowWrapper(F &&f) : func(std::move(f)) {}
+  R operator()(Args... args) const {
+    return ValueOrThrow(func(std::forward<Args>(args)...));
+  }
+  F func;
+};
+template <typename R, typename C, typename... Args, typename F>
+struct ValueOrThrowWrapper<absl::StatusOr<R> (C::*)(Args...) const, F> {
+  explicit ValueOrThrowWrapper(F &&f) : func(std::move(f)) {}
+  R operator()(Args... args) const {
+    return ValueOrThrow(func(std::forward<Args>(args)...));
+  }
+  F func;
+};
+
+// For unbound nonstatic member functions, non-const and const versions.
+// `ptmf` stands for "pointer to member function".
+template <typename R, typename C, typename... Args>
+struct ValueOrThrowWrapper<absl::StatusOr<R>(Args...), C> {
+  explicit ValueOrThrowWrapper(absl::StatusOr<R> (C::*ptmf)(Args...))
+      : ptmf(ptmf) {}
+  R operator()(C &instance, Args... args) {
+    return ValueOrThrow((instance.*ptmf)(std::forward<Args>(args)...));
+  }
+  absl::StatusOr<R> (C::*ptmf)(Args...);
+};
+template <typename R, typename C, typename... Args>
+struct ValueOrThrowWrapper<absl::StatusOr<R>(Args...) const, C> {
+  explicit ValueOrThrowWrapper(absl::StatusOr<R> (C::*ptmf)(Args...) const)
+      : ptmf(ptmf) {}
+  R operator()(const C &instance, Args... args) const {
+    return ValueOrThrow((instance.*ptmf)(std::forward<Args>(args)...));
+  }
+  absl::StatusOr<R> (C::*ptmf)(Args...) const;
+};
+
+}  // namespace jax
+
+#endif  // JAXLIB_ABSL_STATUS_CASTERS_H_