Upstream open-source contributions to foundational machine learning libraries and deep-tech ecosystems (including PyTorch, TensorFlow, and AMD ROCm). When hitting development walls or boredom during local AI modeling, I actively transition to diagnosing issue trackers, tracking framework bugs, writing optimization patches, and submitting Pull Requests directly to major AI infrastructure repos to support the global developer stack.
I have made a PR which fixes an issue in jax-ml/jax repo.
This PR resolves an optimization trajectory divergence bug (detailed in #120986) that occurs when utilizing tf.nn.compute_average_loss inside a distributed multi-worker training loop (MultiWorkerMirroredStrategy) wrapped with XLA compilation (jit_compile=True).
The optimization trajectory splits between the non-XLA (reference) and XLA-compiled execution arms because of how active strategy contexts are evaluated across the function tracing boundary.
Historically, compute_average_loss relied exclusively on distribute_lib.get_strategy().num_replicas_in_sync to infer the global batch size scaling constant. This utility queries the thread-local strategy_stack maintained by the eager runtime context. However, during the graph tracing initialization phase of tf.function(jit_compile=True), this eager stack handle can become isolated or masked inside the polymorphic function builder thread space. Furthermore, if a trace is evaluated prematurely or during initialization steps, the stack can appear empty to the compilation tracer frame.
When this handle is masked or evaluates to empty, compute_average_loss silently falls back to a default global_batch_size = 1. Under XLA, this value is permanently constant-folded and baked directly into the HLO binary graph. During actual parallel execution inside strategy.run(), the underlying math executes math_ops.reduce_sum(per_example_loss) / 1 instead of dividing by the correct num_replicas_in_sync (e.g., 2). This permanently doubles the gradient magnitude on the XLA arm, causing stateful optimizers like Adam to immediately diverge.
_get_num_replicas_in_sync() inside tensorflow/python/ops/nn_impl_distribute.py. This helper modifies the strategy discovery topology to prioritize checking the lightweight, per-thread active environment handle distribute_lib.get_replica_context(). This context is guaranteed to be populated inside strategy.run() and cleanly survives the Python frame switches involved in XLA function tracing.int to catch structural strategy bugs prior to HLO constant generation.compute_average_loss to route dynamic scaling factor discovery through the new _get_num_replicas_in_sync() pipeline.TestGetNumReplicasInSync and TestComputeAverageLossXLAContextMasking) directly to tensorflow/python/ops/nn_loss_scaling_utilities_test.py to formally mock context-masking behavior and assert proper scalar output reductions under isolated tracing simulation environments.py_compile checks.
I have force-pushed an updated, compile-safe implementation to address the logical and architectural edge cases across both the CPU and CUDA backends.
if constexpr (std::is_integral_v<scalar_t>) branching to decouple the type dispatch loops. This completely isolates the NaN-preservation logic to floating-point types (float, double, Half, BFloat16). As a result, integral type paths completely elide the isnan guards at compile time, eliminating template instantiation errors and preserving the highly optimized native fast paths.sign_kernel_cuda, the host-side at::_isnan utility inside the GPU_LAMBDA has been replaced with the device-safe qualified c10::isnan intrinsic. This resolves the NVCC __host__ to __device__ namespace violation completely.spin lint and lintrunner -a pass with zero code style or formatting errors, and local target compilation succeeds perfectly.The PR is now fully production-ready and structurally sound.
Investigated and resolved a silent miscompilation bug in torch.compile where abs() operations on unsigned integer tensors (such as uint8) routed through the C++ codegen produced incorrect downstream graph results (Fixes #187018).
std::abs({x}) unconditionally for scalar paths. In C++, std::abs lacks an overload for unsigned types, triggering implicit signed integer promotion that broke critical downstream unsigned wrap-around semantics during negation and reduction.CppOverrides.abs() within torch/_inductor/codegen/cpp.py to handle types explicitly. Unsigned integers now bypass C++ math entirely as an identity operation (x), signed integers utilize a typed ternary bypass to avoid implicit promotion, and floating-point paths preserve the standard std::abs({x}) fallback.test_uint8_abs_codegen_issue_187018) within test/inductor/test_torchinductor.py under CpuTests. Verified that eager and compiled outputs now align flawlessly with strict parity.
This patch resolves a critical numerical divergence inside the PyTorch core library regarding how torch.sign handles NaN (Not a Number) floating-point elements. By replacing relational comparison dependencies with an explicit NaN-interception architecture, this modification ensures strict adherence to modern IEEE 754 floating-point standards and unifies execution parity across both CPU (SIMD Vectorized) and GPU (CUDA Threaded) runtimes.
Previously, both the CPU and CUDA hardware backend kernels computed the numerical sign of a tensor element using raw relational comparison arithmetic expressions equivalent to:
$$\text{Result} = (0 < a) - (a < 0)$$
According to the IEEE 754 specification, any relational comparison operations evaluated directly against a NaN element must return false. Because of this constraint, when a NaN entered the execution block:
This configuration silently swallowed the arithmetic NaN state, transforming an undefined input value into a normal, valid 0 tracking tensor, which broke expected downstream behavior in mathematical gradients and loss validation.
To fix this without degrading processing speeds, the implementation handles the problem at the hardware compilation level for both devices:
Advanced CPU Kernel Path (UnaryOpsKernel.cpp):
Scalar Fallback: Integrated an immediate type-safe intercept layer using at::_isnan(a). If true, the branch bypasses relational arithmetic entirely and returns the original NaN element.
SIMD Vectorization: Packaged hardware registers to calculate masks using self_vec.isnan(). It applies a precise bitwise rewrite step via Vectorized<scalar_t>::blendv to preserve original uncorrupted data lanes right before register return.
Parallel GPU Kernel Path (UnarySignKernels.cu):
Upgraded the global device execution lambda ([]GPU_LAMBDA) to intercept thread blocks via at::_isnan(a) before the variable drops down into the low-level c10::signum(a) macro block.
