Repository navigation
fix(tf2xla): register EuclideanNorm reduction OpKernel - #126936
Conversation
dmiltr3
left a comment
dmiltr3
left a comment
There was a problem hiding this comment.
Thanks for restoring the EuclideanNormOp kernel registration!
The GitHub CI tests (e.g., Linux x86 CUDA) have failed during compilation. Our internal local validation shows that the compilation fails with error: no member named 'MaybeConjugate' in namespace 'xla'.
To fix this, could you please add the following includes to tensorflow/compiler/tf2xla/kernels/reduction_ops.cc:
#include "xla/hlo/builder/lib/math.h"
#include "xla/primitive_util.h"You might also need to update the Bazel .BUILD dependencies for reduction_ops to include //tensorflow/compiler/xla/hlo/builder/lib:math to resolve linker/compilation errors.
Once you have pushed those fixes, we'll rerun tests and get this merged!
github-project-automation
Bot
moved this from Assigned Reviewer
to Reviewer Requested Changes
in PR Queue
|
Thanks @dmiltr3. Pushed two commits:
CI should rebuild on the new tip. |
github-project-automation
Bot
moved this from Reviewer Requested Changes
to Approved by Reviewer
in PR Queue
nithyak0204
removed
awaiting review
dmiltr3
left a comment
dmiltr3
left a comment
There was a problem hiding this comment.
Thank you for your contribution!
While the automated test suite passed, the PR cannot be merged yet due to an upstream merge conflict in tensorflow/compiler/jit/mark_for_compilation_pass.cc.
Root Cause
PR #126935 was recently merged into upstream master, which added BesselI0 and BesselI1 to GetAllowlistTable() in tensorflow/compiler/jit/mark_for_compilation_pass.cc. This causes a conflict with the changes on this branch that add EuclideanNorm.
Required Fix
Please rebase your branch on the latest upstream master and ensure both changes are retained in tensorflow/compiler/jit/mark_for_compilation_pass.cc:
git fetch upstream
git rebase upstream/masterIn tensorflow/compiler/jit/mark_for_compilation_pass.cc, verify that GetAllowlistTable() retains both BesselI0/BesselI1 in the unary ops list and EuclideanNorm in the reduction ops list:
// Unary / pointwise allowlist (from #126935):
"Lgamma", "Digamma", "BesselI0", "BesselI1",
// Reduction allowlist:
{"All", "Any", "Min", "Max", "Mean", "Prod", "Sum", "EuclideanNorm"}},Once rebased, please push the updated branch to GitHub so our automated synchronization pipeline can cleanly import and complete the merge.
github-project-automation
Bot
moved this from Approved by Reviewer
to Reviewer Requested Changes
in PR Queue
Compilation failed with no member MaybeConjugate in namespace xla. Includes were present; reduction_ops BUILD lacked @xla//xla/hlo/builder/lib:math.
Newly registered XLA EuclideanNorm must appear in GetAllowlistTable. Add it to the RED category with Sum/Prod/Mean.
5ac4a15 to
c51805a
Compare
|
Rebased |
dmiltr3
left a comment
dmiltr3
left a comment
There was a problem hiding this comment.
The C++ changes introduced in this PR look good and provide the missing includes to correctly build reduction_ops targets. I successfully verified the code in Google's internal TAP CI, and it builds properly out-of-the-box.
I noticed a couple of our longer-running CI targets timed out during my verify (//learning/brain/google/client:remote_session_test_gpu and //third_party/tensorflow/python/data/kernel_tests:repeat_test), but these timeouts appear to be pre-existing infrastructure flakiness totally unrelated to your changes. Therefore I'm approving this PR to go ahead. Thank you for the contribution!
github-project-automation
Bot
moved this from Reviewer Requested Changes
to Approved by Reviewer
in PR Queue
5d2d621
into
tensorflow:master
github-project-automation
Bot
moved this from Approved by Reviewer
to Merged
in PR Queue
Why
In commit 8e01ae8 (August 2020), an inverted commit during an automated import accidentally removed the
EuclideanNormOpkernel registration, thePreprocessInputhook onXlaReductionOp, andtestReduceEuclideanNormfromtf2xla.Since then, compiling any graph containing
tf.math.reduce_euclidean_normunder XLA (jit_compile=True) fails withInvalidArgumentError: Detected unsupported operations when trying to compile graph ... on XLA_CPU_JIT: EuclideanNorm (No registered 'EuclideanNorm' OpKernel for XLA_CPU_JIT devices compatible with node).This patch restores the
PreprocessInputextension point toXlaReductionOpand registers theEuclideanNormreduction OpKernel intf2xla.Scope
tensorflow/compiler/tf2xla/kernels/reduction_ops.h, addedvirtual xla::XlaOp PreprocessInput(xla::XlaBuilder* builder, const xla::XlaOp& data)with default pass-through behavior.tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc, implemented defaultPreprocessInputand called it inXlaReductionOp::Compileprior to the reduction computation.tensorflow/compiler/tf2xla/kernels/reduction_ops.cc, implemented and registeredEuclideanNormOpusingxla::Mul(data, xla::MaybeConjugate(data, true))inPreprocessInput,xla::AddinBuildReducer, andxla::SqrtinBuildFinalizer.tensorflow/compiler/tests/reduce_ops_test.py, restoredtestReduceEuclideanNormcoveringfloat32andfloat64.Blast Radius
This change affects the XLA lowering path for
EuclideanNorm. Standard non-XLA execution paths are unaffected. Existing reductions (Sum,Prod,Min,Max,Mean,All,Any) inherit the defaultPreprocessInputpass-through with no behavioral changes.Verification
InvalidArgumentError: No registered 'EuclideanNorm' OpKernel for XLA_CPU_JIT devices.tensorflow/compiler/tests/reduce_ops_test.pyvalidates reduction across axes against a NumPy reference implementation forfloat32andfloat64.