Check shape, ndim and dtype of tensor/ndarray of input of function
Project description
ArrayContract
from arraycontract import shape, _
import torch
@shape(x=(_, 'N'), y=('N', _))
def matrix_dot(x, y):
return x @ y
matrix_dot(torch.rand(3,4), torch.rand(4,5)) # OK
matrix_dot(torch.rand(3,4), torch.rand(3,5)) # raise AssertionError
from arraycontract import shape, _
import torch
from torch import nn
linear = nn.Linear(3, 4)
@shape((..., 3))
def forward_linear(x):
"""
requires x.shape[-1] == 3
"""
return linear(x)
forward_linear(torch.rand(4,5,3)) # OK
forward_linear(torch.rand(4,4)) # raise AssertionError
from arraycontract import dtype
from arraycontract import ndim
import torch
@ndim(x=3, y=4)
def ndim_contract(x, y):
print("requires x.ndim == 3 and y.ndim == 4")
@dtype(x=torch.long)
def dtype_contract(x):
print("requires x.dtype == torch.long")
from arraycontract import Trigger
from arraycontract import dtype
import torch
Trigger.dtype_check_trigger = False
@dtype(x=torch.long)
def dtype_contract(x):
print("not requires x.dtype == torch.long")
dtype_contract(torch.rand(3, 4).float()) # OK
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
arraycontract-0.0.2.tar.gz
(3.8 kB
view hashes)
Built Distribution
Close
Hashes for arraycontract-0.0.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 65bdc9cf8f495aee205d101f2974ab8153bb7e385f4ae23e23572fc76c78a6ab |
|
MD5 | b4479903c980c66f94e4b413c7ecb8bd |
|
BLAKE2b-256 | ba84266438a9e442809c93f95f1d0ef1bd9ab4274b5862287b416ce98021c382 |