`UniformFloatingPointDistribution` incorrect behaviour at infinite bounds

[joelberkeley]

Copied from a TensorFlow issue I raised and closed

Issue Type

Bug

Source

source (I actually use this binary but that isn't compiled by Google))

Tensorflow Version

2.8

Custom Code

Yes

OS Platform and Distribution

Ubuntu 20.04

Mobile device

n/a

Python version

n/a

Bazel version

2.4.1

GCC/Compiler version

Unknown, between 7.5 and 9.3

CUDA/cuDNN version

CUDA Version: 11.6

GPU model and memory

NVIDIA GeForce GTX 1070

Current Behaviour?

`UniformFloatingPointDistribution` produces the following samples for the following bounds

1) -inf 0 -> nan
2) 0 +inf -> inf
3) -inf -inf -> nan
4) +inf +inf -> nan

I believe 1) is incorrect and inconsistent with 2), which is correct. I believe 3) and 4) should be -inf and +inf respectively, since any sample between one +inf and another +inf will be +inf, and since there's no way to specify _different_ +infs (so that bounds are different), I think it makes sense to assume the bounds of +inf and +inf are different, and the same for -inf and -inf.

Standalone code to reproduce the issue

#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/client/lib/constants.h"
#include "tensorflow/compiler/xla/client/lib/prng.h"
#include "tensorflow/compiler/xla/shape.h"
#include "tensorflow/compiler/xla/client/local_client.h"
#include "tensorflow/compiler/xla/client/client_library.h"
#include "tensorflow/core/common_runtime/gpu/gpu_init.h"

void Test() {
    xla::XlaBuilder builder("");
    auto posinf = xla::MaxValue(&builder, xla::F64);
    auto neginf = xla::MinValue(&builder, xla::F64);
    auto zero = xla::ConstantR0<double>(&builder, 0.0);
    auto key = xla::ConstantR0<uint64_t>(&builder, 0);
    auto state = xla::ConstantR0<uint64_t>(&builder, {0});
    auto shape = xla::ShapeUtil::MakeShape(xla::F64, {});

    auto sample = xla::UniformFloatingPointDistribution(
        key, state, xla::ThreeFryBitGenerator, neginf, 0, shape // replace bounds as appropriate
    ).value;

    auto computation = builder.Build(sample).ConsumeValueOrDie();
    auto res =
        xla::ClientLibrary::GetOrCreateLocalClient(tensorflow::GPUMachineManager())  // I'm also seeing this on CPU
        .ConsumeValueOrDie()
        ->ExecuteAndTransfer(computation, {})
        .ConsumeValueOrDie()
        .ToString();

    std::cout << res << std::endl;
}

Relevant log output

f64[] -nan
f64[] inf
f64[] -nan
f64[] -nan

for each test case 1-4

[cheshire]

@burmako Seems like something you would want to take on?

[burmako]

@cheshire @tpopp Thank you for looping me in! Within the StableHLO team, we're focused on defining the semantics of the StableHLO opset as well as implementing it in the StableHLO dialect. This work also includes converters to/from MHLO and HLO.

Ensuring the semantics of the HLO opset and/or supporting the client HLO APIs like UniformFloatingPointDistribution and UniformIntDistribution is not something that's currently on our radar, so I think we'll need to ask someone else for guidance. Apologies for not being more helpful!

[tpopp]

The expected range is [a,b) in all distribution functions that I've seen on real values, so I think points 1 and 2 are bugs. With this, the user would be expected to handle [x,x) if they expect that to return x, as that is not the common behavior. Please correct me though if you've seen this behavior used/desired commonly.

[joelberkeley]

I guess this is a problem that inf compares equal to inf even though it's very questionable mathematically. I was thinking that [inf, inf) is for different infinities, which they can be but aren't necessarily. OK I'll go back to the drawing board