`lattice_rans` compression fails with cudaErrorInvalidValue on bfloat16 tensors, silently falls back to uncompressed storage

The failure is the cudaErrorInvalidValue raised inside the lattice_rans encoder when given a (256, 256) bfloat16 tensor; the fallback path that stores the tensor uncompressed also needs to be located. Start by reproducing with the snippet in the issue on a CUDA device and tracing where the fallback warning is emitted. Done means either bfloat16 encoding succeeds at about 3.5 bpp, or compress() raises a clear error instead of falling back silently.

Description

When compressing a torch.bfloat16 tensor on CUDA-compatible accelerators (T-Head SAIL ZW-M890P) using compress_method='lattice_rans' (the default), the compression fails internally with cudaErrorInvalidValue: invalid argument. Instead of raising an error, entropack silently falls back to storing the tensor uncompressed at ~16 bpp, despite requesting target_bpp=3.5.

This makes it appear as if compression succeeded (the API returns normally), but the output is actually uncompressed data with a misleading actual_bpp value.

Reproduction

import entropack as ep
import torch

tensor = (torch.randn(256, 256, device="cuda") * 0.02).to(torch.bfloat16)
config = ep.EntroPackLossyConfig(target_bpp=3.5)

compressed = ep.compress(tensor, config)
restored = ep.decompress(compressed, config)

print(f"Target: {config.target_bpp:.2f} bits per parameter")
print(f"Stored: {compressed.actual_bpp:.2f} bits per parameter")
print(restored.shape, restored.dtype)

Actual Output

lattice_rans could not encode a (256, 256) torch.bfloat16 tensor, storing it uncompressed: cudaErrorInvalidValue: invalid argument

CompressionFallbackWarning: entropack stored a (256, 256) torch.bfloat16 tensor uncompressed, at 16 bits per parameter: compress_method='lattice_rans' failed with CUDARuntimeError: cudaErrorInvalidValue: invalid argument

Target: 3.50 bits per parameter
Stored: 16.04 bits per parameter
torch.Size([256, 256]) torch.bfloat16

Expected Behavior

Either:

1. The lattice_rans method should correctly handle bfloat16 tensors of shape (256, 256), achieving the target bitrate (~3.5 bpp).
2. Or, if bfloat16 is not supported by lattice_rans, the library should raise a clear error at compress() time rather than silently falling back to uncompressed storage.

A silent fallback to 16 bpp (uncompressed) when the user explicitly requested 3.5 bpp is misleading — the CompressionFallbackWarning is easy to miss, and the returned actual_bpp of 16.04 vs the target of 3.5 represents a 4.6x difference in storage.

Environment

• entropack version: 0.1.5
• Python version: 3.12
• PyTorch version: 2.11.0+v0.1.0.ppu2.1.0.oe
• CUDA version: 13.0.0
• GPU: ZW-M890P
• OS: Ubuntu 24.04

Additional Context

• The input tensor values are small (* 0.02), which should be favorable for compression — this rules out the possibility that the data is simply incompressible.
• The assert torch.allclose(tensor, restored, atol=1e-3) passes trivially because the data was stored uncompressed.