From 110e102e5ef8bfde3e89ea5dbd57dccf7add1521 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 3 Jun 2025 12:59:43 -0700 Subject: [PATCH 1/2] Break JAX extensions --- transformer_engine/jax/csrc/extensions.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index acc9c1dcb1..bdd102101a 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -7,6 +7,8 @@ #ifndef TRANSFORMER_ENGINE_JAX_CSRC_FP8_MODULES_H_ #define TRANSFORMER_ENGINE_JAX_CSRC_FP8_MODULES_H_ +break the code here + #include #include #include From 932f85df361d8290bc2af0a0f87750f6b750fb5c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 3 Jun 2025 20:00:26 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/jax/csrc/extensions.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index bdd102101a..4d71da4202 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -33,8 +33,8 @@ break the code here #include "extensions/utils.h" #include "transformer_engine/activation.h" -// ENUM_ATTR and DICT_ATTR recoding need to be registered in the global namespace -XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Scaling_Mode); + // ENUM_ATTR and DICT_ATTR recoding need to be registered in the global namespace + XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Scaling_Mode); namespace transformer_engine { namespace jax {