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
from py_jaxtyping import PyArray
from jaxtyping import Int
import numpy as np
class Sample(BaseModel):
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.0.tar.gz
(2.5 kB
view hashes)
Built Distribution
Close
Hashes for py_jaxtyping-0.1.0-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | aba2f5559959d7ca6ebc54e4a3f768a99449514e55cf824b6ab55f4821437b3e |
|
MD5 | f566de0c413300970058f7bd317e96a3 |
|
BLAKE2b-256 | 3996576c74d8b3bb1d115862ff980da53ee947462e272c139bd735454695dc7a |