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.3.tar.gz
(2.8 kB
view hashes)
Built Distribution
Close
Hashes for py_jaxtyping-0.1.3-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 4e00f5e1089995de1851e76bdadfed4b3da4331b42ec0b8817c1ff53e877d59a |
|
MD5 | 98072a035085c7f156d08d098c2bf82a |
|
BLAKE2b-256 | 0233969ba8cf53887770f5422d9e32c81399f4eee69c93a1091edc766e9a05be |