patrick-kidger / patrick-kidger/jaxtyping

Jaxtyping annotations don't work with pyserde

Open
#352 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

question
Dominant language
Python
Stars
1.9k
Forks
96
PR merge metrics
No merged PRs in 30d

Description

I'm using Jaxtyping to make sure that my tensors are in the correct shape. However, I also need to serialise the tensors to an API, for which I'm using pyserde with a custom global serialiser.
If the types of the tensor are plain torch.Tensor the custom serialisers works and I can serialise the classes to JSON or other pyserde supported formats. However, when I annotate the tensors with Jaxtyping, e.g. Float32[torch.Tensor, "4"] the serialiser fails with the error:

serde.compat.SerdeError: Unsupported type: Tensor

Minimal (not) working example, uncomment the field to get the error:

from typing import Annotated, Any
from plum import dispatch
import torch

from jaxtyping import Float32
import serde
from serde.json import to_json


class Serializer:
    @dispatch
    def serialize(self, value: torch.Tensor) -> Any:
        return {
            "__tensor__": True,
            "dtype": str(value.dtype),
            "shape": list(value.shape),
            "data": value.cpu().tolist(),
        }


serde.add_serializer(Serializer())


@serde.serde
class Foo:
    tensor_works: torch.Tensor

    annotated_works: Annotated[torch.Tensor, "4"]

    # jax: Float32[torch.Tensor, "4"]


foo = Foo(
    tensor_works=torch.tensor([1000.0, 2000.0, 3000.0]),
    annotated_works=torch.tensor([0.0, 0.0, 0.0, 1.0], dtype=torch.float32),
    # jax=torch.tensor([100.0, 0.0, 0.0, 0.0], dtype=torch.float32), <- Doesn't work
)

j = to_json(foo)
print(j)

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the minimal example, focusing on serde.add_serializer and the to_json entry point; compare how plain torch.Tensor and the Float32[...] annotation are handled. Done means the provided example serialises the annotated tensor to JSON without the Unsupported type: Tensor error.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
api, backend
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.