Merge pull request #310 from bbopen/feat/0.10-torch-dtypes

feat(codec): torch bfloat16 via exact upcast + source provenance

Author: bbopen

docs/codec-envelopes.md

@@ -78,8 +78,10 @@ The target is `torch.Tensor`. CPU tensors work by default. Schema sketch:
   "codecVersion": 1,
   "encoding": "ndarray",
   "shape": [ ... ],
-  "dtype": "float32",
+  "dtype": "torch.float32",
   "device": "cpu" | "cuda:0" | ...,
+  "sourceDtype": "torch.bfloat16", // optional; original dtype before conversion
+  "sourceDevice": "cuda:0",       // optional; original device before CPU copy
 
   // Nested ndarray envelope:
   "value": {
@@ -93,6 +95,9 @@ The target is `torch.Tensor`. CPU tensors work by default. Schema sketch:
 }
 ```
 
+`torch.bfloat16` tensors are upcast exactly to `torch.float32` for transport,
+with the original dtype recorded in `sourceDtype`.
+
 Set `TYWRAP_TORCH_ALLOW_COPY=1` only when the CPU transfer or contiguous copy
 is acceptable for the call.
 

docs/public/llms-full.txt

@@ -2239,6 +2239,10 @@ interface TorchTensor {
   shape: number[];
   dtype?: string;
   device?: string;
+  /** Original dtype when the bridge upcast the tensor (e.g. torch.bfloat16). */
+  sourceDtype?: string;
+  /** Original device when TYWRAP_TORCH_ALLOW_COPY=1 moved the tensor to CPU. */
+  sourceDevice?: string;
 }
 ```
 
@@ -2401,8 +2405,10 @@ The target is `torch.Tensor`. CPU tensors work by default. Schema sketch:
   "codecVersion": 1,
   "encoding": "ndarray",
   "shape": [ ... ],
-  "dtype": "float32",
+  "dtype": "torch.float32",
   "device": "cpu" | "cuda:0" | ...,
+  "sourceDtype": "torch.bfloat16", // optional; original dtype before conversion
+  "sourceDevice": "cuda:0",       // optional; original device before CPU copy
 
   // Nested ndarray envelope:
   "value": {
@@ -2416,6 +2422,9 @@ The target is `torch.Tensor`. CPU tensors work by default. Schema sketch:
 }
 ```
 
+`torch.bfloat16` tensors are upcast exactly to `torch.float32` for transport,
+with the original dtype recorded in `sourceDtype`.
+
 Set `TYWRAP_TORCH_ALLOW_COPY=1` only when the CPU transfer or contiguous copy
 is acceptable for the call.
 

docs/reference/type-mapping.md

@@ -126,6 +126,10 @@ interface TorchTensor {
   shape: number[];
   dtype?: string;
   device?: string;
+  /** Original dtype when the bridge upcast the tensor (e.g. torch.bfloat16). */
+  sourceDtype?: string;
+  /** Original device when TYWRAP_TORCH_ALLOW_COPY=1 moved the tensor to CPU. */
+  sourceDevice?: string;
 }
 ```
 

runtime/tywrap_bridge_core.py

@@ -656,6 +656,8 @@ def serialize_torch_tensor(obj, *, force_json_markers, torch_allow_copy=False):
     import torch  # already importable: is_torch_tensor() gated the dispatch
 
     tensor = obj.detach()
+    source_device = None
+    source_dtype = None
 
     # Sparse tensors (COO/CSR/CSC/BSR/BSC -> any non-strided layout) have no dense
     # numpy representation without a densify step, which is not the round-trip this
@@ -703,19 +705,23 @@ def serialize_torch_tensor(obj, *, force_json_markers, torch_allow_copy=False):
             raise RuntimeError(
                 'Torch tensor is on a non-CPU device; set TYWRAP_TORCH_ALLOW_COPY=1 to allow CPU transfer'
             )
+        source_device = str(tensor.device)
         tensor = tensor.to('cpu')
     if hasattr(tensor, 'is_contiguous') and not tensor.is_contiguous():
         if not torch_allow_copy:
             raise RuntimeError(
                 'Torch tensor is not contiguous; set TYWRAP_TORCH_ALLOW_COPY=1 to allow contiguous copy'
             )
         tensor = tensor.contiguous()
+    if tensor.dtype == torch.bfloat16:
+        source_dtype = str(tensor.dtype)
+        tensor = tensor.float()
     try:
         arr = tensor.numpy()
     except Exception as exc:
         raise RuntimeError('Failed to convert torch.Tensor to numpy') from exc
 
-    return {
+    envelope = {
         '__tywrap__': 'torch.tensor',
         'codecVersion': CODEC_VERSION,
         'encoding': 'ndarray',
@@ -724,6 +730,11 @@ def serialize_torch_tensor(obj, *, force_json_markers, torch_allow_copy=False):
         'dtype': str(tensor.dtype),
         'device': str(tensor.device),
     }
+    if source_dtype is not None:
+        envelope['sourceDtype'] = source_dtype
+    if source_device is not None:
+        envelope['sourceDevice'] = source_device
+    return envelope
 
 
 def serialize_sklearn_estimator(obj):

src/runtime/pyodide-bootstrap-core.generated.ts

@@ -33,6 +33,8 @@ export interface TorchTensor {
   shape?: readonly number[];
   dtype?: string;
   device?: string;
+  sourceDtype?: string;
+  sourceDevice?: string;
 }
 
 export interface SklearnEstimator {
@@ -93,6 +95,8 @@ export type ValueEnvelope =
       readonly shape?: readonly number[];
       readonly dtype?: string;
       readonly device?: string;
+      readonly sourceDtype?: string;
+      readonly sourceDevice?: string;
     }
   | {
       readonly __tywrap__: 'sklearn.estimator';
@@ -1013,18 +1017,47 @@ const decodeTorchTensorEnvelope: EnvelopeHandler = <T>(
     );
   }
   const device = typeof deviceValue === 'string' ? deviceValue : undefined;
+  const sourceDtypeValue = value.sourceDtype;
+  if (
+    sourceDtypeValue !== undefined &&
+    (typeof sourceDtypeValue !== 'string' || sourceDtypeValue.length === 0)
+  ) {
+    throw new Error(
+      'Invalid torch.tensor envelope: sourceDtype must be a non-empty string when provided'
+    );
+  }
+  const sourceDtype = typeof sourceDtypeValue === 'string' ? sourceDtypeValue : undefined;
+  const sourceDeviceValue = value.sourceDevice;
+  if (
+    sourceDeviceValue !== undefined &&
+    (typeof sourceDeviceValue !== 'string' || sourceDeviceValue.length === 0)
+  ) {
+    throw new Error(
+      'Invalid torch.tensor envelope: sourceDevice must be a non-empty string when provided'
+    );
+  }
+  const sourceDevice = typeof sourceDeviceValue === 'string' ? sourceDeviceValue : undefined;
+
+  const tensor = (data: unknown): TorchTensor => ({
+    data,
+    shape,
+    dtype,
+    device,
+    ...(sourceDtype === undefined ? {} : { sourceDtype }),
+    ...(sourceDevice === undefined ? {} : { sourceDevice }),
+  });
 
   const decoded = recurse(nested);
   if (isPromiseLike(decoded)) {
     return decoded.then(data =>
-      tagDecodedShape({ data, shape, dtype, device } satisfies TorchTensor, {
+      tagDecodedShape(tensor(data), {
         marker: 'torch.tensor',
         dims: shape?.length,
         dtype,
       })
     ) as Promise<T | unknown>;
   }
-  return tagDecodedShape({ data: decoded, shape, dtype, device } satisfies TorchTensor, {
+  return tagDecodedShape(tensor(decoded), {
     marker: 'torch.tensor',
     dims: shape?.length,
     dtype,

src/utils/codec.ts

@@ -678,6 +678,62 @@ describe('#234 codec envelope re-validation (JS decoder)', () => {
       ).toThrow(/device must be a non-empty string/);
     });
 
+    it('preserves dtype and device provenance without changing transported metadata', () => {
+      const result = decodeValue({
+        __tywrap__: 'torch.tensor',
+        codecVersion: 1,
+        encoding: 'ndarray',
+        value: {
+          __tywrap__: 'ndarray',
+          codecVersion: 1,
+          encoding: 'json',
+          data: [1, -2.5, 3.140625],
+          shape: [3],
+          dtype: 'float32',
+        },
+        shape: [3],
+        dtype: 'torch.float32',
+        device: 'cpu',
+        sourceDtype: 'torch.bfloat16',
+        sourceDevice: 'cuda:0',
+      });
+
+      expect(result).toMatchObject({
+        data: [1, -2.5, 3.140625],
+        shape: [3],
+        dtype: 'torch.float32',
+        device: 'cpu',
+        sourceDtype: 'torch.bfloat16',
+        sourceDevice: 'cuda:0',
+      });
+    });
+
+    it.each([
+      ['sourceDtype', ''],
+      ['sourceDtype', 123],
+      ['sourceDevice', ''],
+      ['sourceDevice', false],
+    ] as const)('rejects invalid optional %s provenance', (field, invalidValue) => {
+      expect(() =>
+        decodeValue({
+          __tywrap__: 'torch.tensor',
+          codecVersion: 1,
+          encoding: 'ndarray',
+          value: {
+            __tywrap__: 'ndarray',
+            codecVersion: 1,
+            encoding: 'json',
+            data: [1],
+            shape: [1],
+            dtype: 'float32',
+          },
+          shape: [1],
+          dtype: 'torch.float32',
+          [field]: invalidValue,
+        })
+      ).toThrow(new RegExp(`${field} must be a non-empty string when provided`));
+    });
+
     it('rejects a nested value that is not an ndarray envelope', () => {
       expect(() =>
         decodeValue({

test/codec-envelope-validation.test.ts

@@ -253,6 +253,18 @@ def torch_int64() -> object:
     return torch.tensor([1, -2], dtype=torch.int64)
 
 
+def torch_dtype(dtype_name: str) -> object:
+    import torch
+
+    if dtype_name == "uint8":
+        values = [0, 3, 255]
+    elif dtype_name.startswith("int"):
+        values = [1, -2]
+    else:
+        values = [1.5, -2.25]
+    return torch.tensor(values, dtype=getattr(torch, dtype_name))
+
+
 def torch_scalar() -> object:
     import torch
 
@@ -262,7 +274,9 @@ def torch_scalar() -> object:
 def torch_bfloat16() -> object:
     import torch
 
-    return torch.tensor([1.5], dtype=torch.bfloat16)
+    return torch.tensor(
+        [1.0, -2.5, 3.140625, 2**16, 2**32], dtype=torch.bfloat16
+    )
 
 
 def torch_sparse() -> object:
@@ -277,6 +291,12 @@ def torch_quantized() -> object:
     return torch.quantize_per_tensor(torch.tensor([1.0, 2.0]), 0.1, 0, torch.qint8)
 
 
+def torch_meta() -> object:
+    import torch
+
+    return torch.empty(3, device="meta")
+
+
 def torch_complex() -> object:
     import torch
 

test/menagerie/fixtures/library_torture.py

@@ -697,6 +697,22 @@ export const RUNTIME_CATALOGUE: readonly CatalogueRow[] = [
       value: { data: [1, -2], shape: [2], dtype: 'torch.int64', device: 'cpu' },
     },
   }),
+  ...(['float64', 'int8', 'int16', 'int32', 'uint8'] as const).map(dtype => {
+    const data = dtype === 'uint8' ? [0, 3, 255] : dtype.startsWith('int') ? [1, -2] : [1.5, -2.25];
+    return libraryRow({
+      id: `torch-${dtype}`,
+      call: `torch_dtype('${dtype}')`,
+      functionName: 'torch_dtype',
+      args: [dtype],
+      requires: ['torch', 'pyarrow'],
+      status: 'EXPECTED_OK',
+      currentBehavior: `A dense ${dtype} CPU tensor round-trips with its declared dtype.`,
+      expected: {
+        kind: 'match',
+        value: { data, shape: [data.length], dtype: `torch.${dtype}`, device: 'cpu' },
+      },
+    });
+  }),
   libraryRow({
     id: 'torch-scalar',
     call: 'torch_scalar()',
@@ -712,17 +728,26 @@ export const RUNTIME_CATALOGUE: readonly CatalogueRow[] = [
     id: 'torch-bfloat16',
     call: 'torch_bfloat16()',
     requires: ['torch', 'pyarrow'],
-    status: 'LOUD_FAIL',
-    currentBehavior: 'NumPy conversion rejects bfloat16.',
-    expected: error(/Failed to convert torch\.Tensor to numpy/i),
+    status: 'EXPECTED_OK',
+    currentBehavior: 'bfloat16 values upcast exactly to float32 and retain their source dtype.',
+    expected: {
+      kind: 'match',
+      value: {
+        data: [1, -2.5, 3.140625, 2 ** 16, 2 ** 32],
+        shape: [5],
+        dtype: 'torch.float32',
+        sourceDtype: 'torch.bfloat16',
+        device: 'cpu',
+      },
+    },
   }),
   libraryRow({
     id: 'torch-sparse',
     call: 'torch_sparse()',
     requires: ['torch'],
     status: 'LOUD_FAIL',
     currentBehavior: 'Sparse tensor layouts are rejected explicitly.',
-    expected: error(/sparse tensors are not supported/i),
+    expected: error(/convert to a dense CPU tensor explicitly.*tensor\.to_dense/i),
   }),
   libraryRow({
     id: 'torch-quantized',
@@ -731,15 +756,23 @@ export const RUNTIME_CATALOGUE: readonly CatalogueRow[] = [
     featureProbe: 'import torch; print("1" if hasattr(torch, "quantize_per_tensor") else "0")',
     status: 'LOUD_FAIL',
     currentBehavior: 'Quantized tensors are rejected explicitly.',
-    expected: error(/quantized tensors are not supported/i),
+    expected: error(/dequantize explicitly.*tensor\.dequantize/i),
+  }),
+  libraryRow({
+    id: 'torch-meta',
+    call: 'torch_meta()',
+    requires: ['torch'],
+    status: 'LOUD_FAIL',
+    currentBehavior: 'Meta tensors are rejected explicitly.',
+    expected: error(/materialize the tensor on a real device before returning/i),
   }),
   libraryRow({
     id: 'torch-complex',
     call: 'torch_complex()',
     requires: ['torch'],
     status: 'LOUD_FAIL',
     currentBehavior: 'Complex tensors are rejected explicitly.',
-    expected: error(/complex tensors are not supported/i),
+    expected: error(/split into real\/imag components explicitly before returning/i),
   }),
 
   libraryRow({