Pydantic support for Jaxtyping array annotations
Project description
Py Jaxtyping
Pydantic support for Jaxtyping array annotations
Usage
Instead of Int[np.ndarray, 'B 256 64 3']
, do PyArray[Int, int, 'B 256 64 3']
You can use it with jaxtyping
as normal, but also it will:
- Serialize to nested lists
- Validate the correct shape and datatypes from serialized lists
Example
from pydantic import BaseModel, ConfigDict
from py_jaxtyping import PyArray
from jaxtyping import Int
import numpy as np
class Sample(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
img: PyArray[Int, int, "W H 3"]
label: str
Sample.model_validate({
'img': np.ones((256, 64, 3)),
'label': 'car'
})
# checks out!
Sample.model_validate({
'img': np.ones((256, 64, 1)),
'label': 'car'
})
# fails: invalid dims :/
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
py_jaxtyping-0.1.5.tar.gz
(2.8 kB
view hashes)
Built Distribution
Close
Hashes for py_jaxtyping-0.1.5-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 3e3f22dac8bd01081387731e8932ebacd6e0035072b2c53c913ebbef5f50283e |
|
MD5 | 3fc2c4c296c220b5532fc474478dafc6 |
|
BLAKE2b-256 | c7d1b58f144895eba1a333e8592c7bad089872a8796c33eaff977678acbf1121 |