Commit 398fce03 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TensorFlower Gardener
Browse files

Add a Euclidean norm reduction kernel. This implements a fused...

Add a Euclidean norm reduction kernel. This implements a fused sqrt(reduce_sum(x * conj(x))) kernels for CPU (using Eigen) and GPU (using CUB), which is more efficient than the composite implementation at the TF level. It will also be easier to avoid the issue of producing NaNs in the gradient at the origin.

Adds tf.math.reduce_euclidian_norm() Python interface to call the fused reduction kernel directly.

Gradients will be added in a followup change.

PiperOrigin-RevId: 234188431
parent 5113410e
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment