Update clang-format action and add a test

Author: kanglant

.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/action.yaml

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# 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.
@@ -18,34 +18,37 @@ inputs:
     description: 'The clang-format version to use.'
     required: true
     default: "20.1.5"
-  branch_name:
-    description: 'The repository branch used for fetching comparisons.'
-    required: true
-    default: 'main'
 
 runs:
   using: "composite"
   steps:
   - name: "Checking out repository"
     uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-  - name: "Fetch HEAD of ${{ inputs.branch_name }} branch"
-    shell: bash
-    run: |
-      # Silence safe.directory warnings when a job is running within a container 
-      /usr/bin/git config --global --add safe.directory '*'
-      /usr/bin/git fetch origin ${{ inputs.branch_name }} --depth=1
+    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: |
-      REPO_CLANG_FORMAT=".clang-format"
-      ACTION_DEFAULT_CLANG_FORMAT="${{ github.action_path }}/.clang-format.default"
+      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.
-        uvx clang-format==${{ inputs.clang_format_version }} --dry-run --Werror --verbose $(git diff --name-only origin/main HEAD -- '*.cc' '*.h')
+        $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."
-        uvx clang-format==${{ inputs.clang_format_version }} -style=file:$ACTION_DEFAULT_CLANG_FORMAT --dry-run --Werror --verbose $(git diff --name-only origin/main HEAD -- '*.cc' '*.h')
+        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_