pydantic-tensor

Support parsing, validation, and serialization of common tensors (np.ndarray, torch.Tensor, tensorflow.Tensor, jax.Array) for Pydantic.

Installation

pip install pydantic-tensor

Usage

Validation

from typing import Annotated, Any, Literal

import numpy as np
import tensorflow as tf
import torch
from pydantic import BaseModel, Field

from pydantic_tensor import Tensor

# allow only integers greater equal than 2 and less equal than 3
DimType = Annotated[int, Field(ge=2, le=3)]


class Model(BaseModel):
    #              tensor type                          shape                    dtype
    tensor: Tensor[torch.Tensor | np.ndarray[Any, Any], tuple[DimType, DimType], Literal["int32", "int64"]]


parsed = Model.model_validate({"tensor": np.ones((2, 2), dtype="int32")})
# access the parsed tensor via the "value" property
parsed.tensor.value

# invalid shapes
Model.model_validate({"tensor": np.ones((1, 1), dtype="int32")})
Model.model_validate({"tensor": np.ones((4, 4), dtype="int32")})
Model.model_validate({"tensor": np.ones(2, dtype="int32")})
Model.model_validate({"tensor": np.ones((2, 2, 2), dtype="int32")})

# invalid dtype
Model.model_validate({"tensor": np.ones((2, 2), dtype="float32")})

# successfully validate np.ndarray
Model.model_validate({"tensor": np.ones((2, 2), dtype="int32")})
# convert tf.Tensor to torch.Tensor
Model.model_validate({"tensor": tf.ones((2, 2), dtype=tf.int32)})

Parsing

The JSON representation of the tensor contains the:

• binary data of the tensor in little-endian format encoded in Base64
• shape of the tensor
• datatype of the tensor

from typing import Any

import numpy as np
from pydantic import BaseModel

from pydantic_tensor import Tensor


class Model(BaseModel):
    tensor: Tensor[Any, Any, Any]


parsed = Model.model_validate({"tensor": np.ones((2, 2), dtype="float32")})
# parse to JSON: {"tensor":{"shape":[2,2],"dtype":"float32","data":"AACAPwAAgD8AAIA/AACAPw=="}}
json_dump = parsed.model_dump_json()
# parse back to tensor: array([[1., 1.], [1., 1.]], dtype=float32)
Model.model_validate_json(json_dump).tensor.value

DType Collections

Types Int, UInt, Float, Complex, BFloat from pydantic_tensor.types are unions of dtypes according to their names. For Example Int is defined as Literal["int8", "int16", "int32", "int64"].

from typing import Any

import numpy as np
from pydantic import BaseModel

from pydantic_tensor import Tensor
from pydantic_tensor.types import Int


class Model(BaseModel):
    tensor: Tensor[Any, Any, Int]


for dtype in ["int8", "int16", "int32", "int64"]:
    Model.model_validate({"tensor": np.ones((2, 2), dtype=dtype)})  # success

Model.model_validate({"tensor": np.ones((2, 2), dtype="float32")})  # failure

Lazy Tensors

Use JaxArray, NumpyNDArray, TensorflowTensor, TorchTensor for lazy versions of tensors types. They only handle tensors when their equivalent libraries (jax, numpy, tensorflow, torch) are imported somewhere else in the program.

from typing import Any

import numpy as np
from pydantic import BaseModel

from pydantic_tensor import Tensor
from pydantic_tensor.backend.torch import TorchTensor


class Model(BaseModel):
    tensor: Tensor[TorchTensor, Any, Any]


Model.model_validate({"tensor": np.ones((2, 2), dtype="float32")})  # failure

import torch

Model.model_validate({"tensor": np.ones((2, 2), dtype="float32")})  # success

Development

Install pre-commit hooks

pre-commit install

Lint

hatch run lint:all

Test

hatch run test:test

Check spelling

hatch run spell