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.4.tar.gz
(2.8 kB
view hashes)
Built Distribution
Close
Hashes for py_jaxtyping-0.1.4-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 93a980ea58201bff20cdb4301c9a5c7fed4e5d5558e1fe3b8df25ce2c81b3519 |
|
MD5 | 91d7b62f7caca32b14e67c03935efda4 |
|
BLAKE2b-256 | 5b6512f1482a82b627c7ce8c7044af294f7a673dfd6d1294294e60ceade59093 |