ConvertElementType bool-to-int returns incorrect result for nonstandard booleans

[jakevdp]

First reported at jax-ml/jax#34751.

StableHLO convert spec says:

For boolean-to-any-supported-type conversions, the value false is converted to zero, and the value true is converted to one.

However, for non-standard boolean representations, XLA does not conform to this when converting to integer:

>>> import jax
>>> import numpy as np
>>> nonstandard_bool = jax.numpy.asarray(np.frombuffer(b'\x00\x01\xff', dtype=bool))
>>> print(nonstandard_bool)
[False  True  True]

>>> print(jax.lax.convert_element_type(nonstandard_bool, 'int32'))
[  0   1 255]

# HLO of the above:
>>> print(jax.jit(lambda x: jax.lax.convert_element_type(x, 'int32')).lower(nonstandard_bool).as_text())
module @jit__lambda attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<3xi1>) -> (tensor<3xi32> {jax.result_info = "result"}) {
    %0 = stablehlo.convert %arg0 : (tensor<3xi1>) -> tensor<3xi32>
    return %0 : tensor<3xi32>
  }
}

[akuegel]

Fixing this for CPU/GPU fusion emitters is relatively straightforward, but doing this for the Interpreter seems difficult. I tried adding the following logic to ConvertBetweenNativeTypes method:

xla/xla/literal.cc

Line 1730 in 49544cc

void ConvertBetweenNativeTypes(absl::Span<const NativeSrcT> src_data,

if constexpr (std::is_same_v<NativeSrcT, bool>) {
  src = static_cast<uint8_t>(src) != 0;
}

But it looks like the C++ compiler optimizes this away somehow. Adding a LOG line to print src in this if statement also leads to a crash.

[akuegel]

Also, I tried this snippet of code:

char val = 255;
bool* pred_val = reinterpret_cast<bool*>(&val);
printf("%d\n", static_cast<int>(*pred_val))

Without -O3 it prints 255 (g++ (Debian 15.2.0-3) 15.2.0) and with -O3 it prints 1. Looks like even C++ has trouble with non-standard booleans.

[akuegel]

@jakevdp It seems that even numpy gets this wrong. After I fixed the issue for CPU/GPU backend, now the testViewScalar10 test form lax_numpy_test is failing in CheckAgainstNumpy:

Mismatched elements: 1 / 1 (100%)
Max absolute difference among violations: 42
Max relative difference among violations: 1.02439024
 ACTUAL: array(1, dtype=int8)
 DESIRED: array(-41, dtype=int8)

[jakevdp]

That test covers view (i.e. BitcastConvertType) which is different from astype (i.e. ConvertElementType).

NumPy is behaving correctly here: view should round trip arbitrary integers through bool without losing bits.

The issue here is that XLA does not support BitcastConvertType for boolean inputs or outputs, and so for JAX's implementation of np.view on booleans, we use ConvertElementType and inadvertently depend on its existing incorrect behavior.

I'm not sure what to do here: if we fix ConvertElementType so that it has correct semantics for bool, then there would be literally no way for XLA to recreate the semantics of np.view unless we also extend BitcastConvertType to work with boolean inputs and outputs. An alternative would be to simply document that jax.numpy.view does not preserve bits when round-tripping through a boolean type.

What do you think?

[jakevdp]

Could we extend XLA to work with boolean types for BitcastConvert? Looking at the stablehlo spec, there doesn't seem to be any prohibition against boolean inputs or outputs: https://openxla.org/stablehlo/spec#bitcast_convert