TensorFlow bcomplex32 pair reproducer

This repository demonstrates a pair-based TensorFlow representation for complex values with BF16 components. In the TensorFlow 2.20 and 2.21 versions tested here, TensorFlow's public DataType exposes complex64 and complex128, but not a public complex<bf16> dtype. The reproducer therefore keeps the real and imaginary components in separate BF16 tensors.

What this demonstrates

For four matching BF16 input planes a_real, a_imag, b_real, and b_imag, the pair path computes:

out_real = bf16(f32(a_real) * f32(b_real) - f32(a_imag) * f32(b_imag))
out_imag = bf16(f32(a_real) * f32(b_imag) + f32(a_imag) * f32(b_real))

Persistent inputs and outputs are BF16 planes, while arithmetic uses FP32 opmath before explicit output quantization.

The same function runs in eager mode and tf.function, preserves first-order autodiff, round-trips through SavedModel, and compiles with XLA when all four plane shapes are fully defined and equal.

Workload

The executable uses a deterministic differentiable frequency-domain signal-and-filter workload with shape (2, 3, 16):

• 2 batch elements
• 3 channels
• 16 frequency bins
• complex signal and filter represented by four BF16 planes
• mean spectral-energy loss
• gradients with respect to all four input planes

A separate oracle branch constructs complex64, multiplies the signal and filter, then quantizes its real and imaginary outputs back to BF16. The pair path must match the quantized oracle exactly for forward values, loss, and all four gradients. The oracle is validation-only and is not part of the pair execution path.

Run locally

Python 3.11 is the CI version.

python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
python tensorflow_bcomplex32_pair_repro.py
python -m pytest -q

Force the standalone evidence run onto CPU:

CUDA_VISIBLE_DEVICES=-1 python tensorflow_bcomplex32_pair_repro.py

GPU evidence requires a TensorFlow GPU build, a compatible visible GPU, and a working CUDA runtime. The final stdout line is machine-readable JSON.

Evidence emitted

The JSON payload reports:

• TensorFlow, NumPy, CUDA, cuDNN, and visible-device metadata
• pair-versus-oracle forward, loss, and gradient errors
• BF16 gradient dtypes for all four input planes
• a GraphDef scan for complex dtype enums and Complex, Conj, Real, or Imag operations
• SavedModel forward equivalence and dynamic shape-mismatch rejection
• XLA output devices and exact HLO scalar-token evidence for bf16, c64, and c128

The GraphDef and HLO checks apply to the pair execution path. They fail if that path contains a hidden complex64 or complex128 fallback.

TensorFlow Assert operations are not used as the XLA shape contract. The eager/graph/SavedModel path supports partially defined shapes with runtime equality checks; the XLA evidence path compiles bcomplex32_mul with a fully defined matching input signature.

Test matrix

Environment	Coverage
GitHub Actions, Python 3.11, TensorFlow 2.20.*	CPU pytest suite and standalone evidence
GitHub Actions, Python 3.11, TensorFlow 2.21.*	CPU pytest suite and standalone evidence
Local compatible TensorFlow GPU environment	GPU XLA execution, HLO tokens, device reporting, and the full pytest suite

Scope and non-goals

This repository is a focused reproducer. It is not:

• a TensorFlow patch or public dtype implementation
• a custom op
• a benchmark or native-acceleration claim
• a packaged or stable API
• a claim that StableHLO/XLA type support automatically creates a TensorFlow public dtype