Skip to content
This repository has been archived by the owner on Aug 15, 2019. It is now read-only.

Commit

Permalink
Made tf.abs underflow-safe on CPU & GPU and added test for it. (#1391)
Browse files Browse the repository at this point in the history
BUG
  • Loading branch information
DirkToewe authored and Nikhil Thorat committed Nov 15, 2018
1 parent 4280e76 commit 5e8ce74
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 5 deletions.
2 changes: 1 addition & 1 deletion src/kernels/backend_cpu.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1122,7 +1122,7 @@ export class MathBackendCPU implements KernelBackend {
for (let i = 0; i < x.size; ++i) {
const real = values[i * 2];
const imag = values[i * 2 + 1];
resultValues[i] = Math.sqrt(real * real + imag * imag);
resultValues[i] = Math.hypot(real, imag);
}
return Tensor.make(x.shape, {values: resultValues}) as T;
}
Expand Down
13 changes: 9 additions & 4 deletions src/kernels/webgl/complex_abs_gpu.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,16 @@ export class ComplexAbsProgram implements GPGPUProgram {
this.outputShape = shape;
this.userCode = `
void main() {
float real = getRealAtOutCoords();
float imag = getImagAtOutCoords();
vec2 v = vec2(real, imag);
float re = abs(getRealAtOutCoords());
float im = abs(getImagAtOutCoords());
float mx = max(re, im);
setOutput(sqrt(dot(v, v)));
// sadly the length function in glsl is not underflow-safe
// (at least not on Intel GPUs). So the safe solution is
// to ensure underflow-safety in all cases.
setOutput(
mx == 0.0 ? 0.0 : mx * length(vec2(1, min(re, im)/mx))
);
}
`;
}
Expand Down
29 changes: 29 additions & 0 deletions src/ops/unary_ops_test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import * as tf from '../index';
import {describeWithFlags} from '../jasmine_util';
import {ALL_ENVS, expectArraysClose, expectNumbersClose} from '../test_util';
import * as util from '../util';
import {ENV} from '../environment';

import * as selu_util from './selu_util';

Expand Down Expand Up @@ -168,6 +169,34 @@ describeWithFlags('abs', ALL_ENVS, () => {
expect(result.shape).toEqual([2, 2, 2]);
});

it('is underflow-safe for complex64', () => {

const floatBits = ENV.backend.floatPrecision();
let small;
switch(floatBits) {
case 32: small = 1e-30; break;
case 16: small = 1e-4; break;
default: throw new Error(
`Test not implemented for ENV.engine.floatPrecision()=${floatBits}.`
);
}

const a = tf.complex(
[small, 0, small, 0],
[small, small, 0, 0]
);
const result = tf.abs(a);
expectArraysClose(
result,
[Math.hypot(small, small),
Math.hypot( 0, small),
Math.hypot(small, 0),
Math.hypot( 0, 0)],
/*tolerance=*/small/100
);
expect(result.shape).toEqual([4]);
});

it('propagates NaNs', () => {
const a = tf.tensor1d([1, -2, 0, 3, -0.1, NaN]);
const result = tf.abs(a);
Expand Down

0 comments on commit 5e8ce74

Please sign in to comment.