feat: support bf16 from pytorch dataset

[eddyxu]

Summary

Support round-trip to use bf16 from PyTorch

Co-authored-by: Claude Opus 4.6 noreply@anthropic.com

[github-actions]

PR Review

P0 Bug: uint64_as_int64 conversion broken for numeric types

The insertion of the elif _is_bfloat16_type(arr.type): branch has accidentally moved the uint64_as_int64 guard into the wrong branch.

Before (base):

        elif (
            pa.types.is_integer(arr.type)
            or pa.types.is_floating(arr.type)
            or pa.types.is_boolean(arr.type)
        ):
            tensor = torch.from_numpy(arr.to_numpy(zero_copy_only=False))

            if uint64_as_int64 and tensor.dtype == torch.uint64:  # ← inside numeric branch ✓
                tensor = tensor.to(torch.int64)

After (this PR):

        elif (
            pa.types.is_integer(arr.type)
            ...
        ):
            tensor = torch.from_numpy(arr.to_numpy(zero_copy_only=False))
        elif _is_bfloat16_type(arr.type):
            tensor = _bf16_to_tensor(arr)

            if uint64_as_int64 and tensor.dtype == torch.uint64:  # ← now inside bf16 branch ✗
                tensor = tensor.to(torch.int64)

The uint64_as_int64 conversion no longer runs for numeric uint64 columns (regression), and is dead code inside the bf16 branch (bf16 tensors will never have torch.uint64 dtype).

Fix: Move the if uint64_as_int64 block back into (or after) the numeric elif, e.g. by dedenting it to run after all branches, or by keeping it inside the numeric branch:

        ):
            tensor = torch.from_numpy(arr.to_numpy(zero_copy_only=False))

            if uint64_as_int64 and tensor.dtype == torch.uint64:
                tensor = tensor.to(torch.int64)
        elif _is_bfloat16_type(arr.type):
            tensor = _bf16_to_tensor(arr)
        elif hf_converter is not None:

Minor: No null handling in _bf16_to_tensor

If a bf16 array contains nulls, _bf16_to_tensor will silently produce garbage values for null slots. This is consistent with the existing numeric path (to_numpy(zero_copy_only=False) also fills nulls with 0), so not blocking — but worth a brief doc comment noting the assumption.

---

Overall the approach (reinterpreting uint16 storage bytes as bfloat16) is sound and the test coverage for the happy path is good.

[westonpace]

Just some nits

[westonpace]

Should we do a sanity check that the data type of storage is a 16-bit type at this point?

[westonpace]

Is the copy here so that the resulting buffer can be mutable?

[westonpace]

Seems like there should be a way to do this without a copy but maybe not.

[eddyxu]

asked claude / codex to do double checks, to make this opportunist

[westonpace]

Kind of a strange comment. I'm not really sure what it means.

[eddyxu]

fixed.