Package: phantom-tensors
Phantom TensorsTensor types with variadic shapes, for any array-based library, that work with both static and runtime type checkersThis project is currently just a rough prototype! Inspired by:phantom-typesThe goal of this project is to let users write tensor-like types with variadic shapes (viaPEP 646) that are:Amendable tostatic type checking (without mypy plugins).E.g., pyright can tell the difference betweenTensor[Batch, Channel]andTensor[Batch, Feature]Useful for performingruntime checks of tensor types and shapes.E.g.,  can validate -- at runtime -- that arrays of typesNDArray[A, B]andNDArray[B, A]indeed have transposed shapes with respect with each other.Compatible withanyarray-based library (numpy, pytorch, xarray, cupy, mygrad, etc.)E.g. A function annotated withx: torch.Tensorcan be passedphantom_tensors.torch.Tensor[N, B, D]. It is trivial to write custom phantom-tensor flavored types for any array-based library.phantom_tensors.parsemakes it easy to declare shaped tensor types in a way that static type checkers understand, and that are validated at runtime:fromtypingimportNewTypeimportnumpyasnpfromphantom_tensorsimportparsefromphantom_tensors.numpyimportNDArrayA=NewType("A",int)B=NewType("B",int)# static: declare that x is of type NDArray[A, B]#         declare that y is of type NDArray[B, A]# runtime: check that shapes (2, 3) and (3, 2)#          match (A, B) and (B, A) pattern across#          tensorsx,y=parse((np.ones((2,3)),NDArray[A,B]),(np.ones((3,2)),NDArray[B,A]),)x# static type checker sees: NDArray[A, B]y# static type checker sees: NDArray[B, A]Passing inconsistent types toparsewill result in a runtime validation error.# Runtime: Raises `ParseError` A=10 and A=2 do not matchz,w=parse((np.ones((10,3)),NDArray[A,B]),(np.ones((3,2)),NDArray[B,A]),)These shaped tensor types are amenable to static type checking:fromtypingimportAnyimportnumpyasnpfromphantom_tensorsimportparsefromphantom_tensors.numpyimportNDArrayfromphantom_tensors.alphabetimportA,B# these are just NewType(..., int) typesdeffunc_on_2d(x:NDArray[Any,Any]):...deffunc_on_3d(x:NDArray[Any,Any,Any]):...deffunc_on_any_arr(x:np.ndarray):...# runtime: ensures shape of arr_3d matches (A, B, A) patternsarr_3d=parse(np.ones((3,5,3)),NDArray[A,B,A])func_on_2d(arr_3d)# static type checker: Error!  # expects 2D arr, got 3Dfunc_on_3d(arr_3d)# static type checker: OKfunc_on_any_arr(arr_3d)# static type checker: OKWrite easy-to-understand interfaces using common dimension names (or make up your own):fromphantom_tensors.torchimportTensorfromphantom_tensors.wordsimportBatch,Embed,Vocabdefembedder(x:Tensor[Batch,Vocab])->Tensor[Batch,Embed]:...Using a runtime type checker, such asbeartypeortypeguard, in conjunction withphantom_tensorsmeans that the typed shape information will be validated at runtime across a function's inputs and outputs, whenever that function is called.fromtypingimportTypeVar,castfromtyping_extensionsimportassert_typeimporttorchastrfrombeartypeimportbeartypefromphantom_tensorsimportdim_binding_scope,parsefromphantom_tensors.torchimportTensorfromphantom_tensors.alphabetimportA,B,CT1=TypeVar("T1")T2=TypeVar("T2")T3=TypeVar("T3")@dim_binding_scope@beartype# <- adds runtime type checking to function's interfacesdefbuggy_matmul(x:Tensor[T1,T2],y:Tensor[T2,T3])->Tensor[T1,T3]:# This is the wrong operation!# Will return shape-(T1, T1) tensor, not (T1, T3)out=x@x.T# We lie to the static type checker to try to get away with itreturncast(Tensor[T1,T3],out)x,y=parse((tr.ones(3,4),Tensor[A,B]),(tr.ones(4,5),Tensor[B,C]),)# At runtime beartype raises:#   Function should return shape-(A, C) but returned shape-(A, A)z=buggy_matmul(x,y)# Runtime validation error!Installationpipinstallphantom-tensorstyping-extensionsis the only strict dependency. Using features fromphantom_tensors.torch(numpy)requires thattorch(numpy) is installed too.Some Lower-Level Details and FeaturesEverything on display here is achieved using relatively minimal hacks (no mypy plugin necessary, no monkeypatching). Presently,torch.Tensorandnumpy.ndarrayare explicitly supported by phantom-tensors, but it is trivial to add support for other array-like classes.Note that mypy does not support PEP 646 yet, but pyright does. You can run pyright on the following examples to see that they do, indeed type-check as expected!Dimension-Binding Contextsphantom_tensors.parsevalidates inputs against types-with-shapes and performstype narrowingso that static type checkers are privy to the newly proven type information about those inputs. It performs inter-tensor shape consistency checks within a "dimension-binding context". Tensor-likes that are parsed simultaneously are automatically checked within a common dimension-binding context.importnumpyasnpimporttorchastrfromphantom_tensorsimportparsefromphantom_tensors.alphabetimportA,B,Cfromphantom_tensors.numpyimportNDArrayfromphantom_tensors.torchimportTensort1,arr,t2=parse(# <- Runtime: enter dimension-binding context(tr.rand(9,2,9),Tensor[B,A,B]),# <-binds A=2 & B=9(np.ones((2,)),NDArray[A]),# <- checks A==2(tr.rand(9),Tensor[B]),# <- checks B==9)# <- Runtime: exit dimension-binding scope#    Statically: casts t1, arr, t2 to shape-typed Tensors# static type checkers now see# t1: Tensor[B, A, B]# arr: NDArray[A]# t2: Tensor[B]w=parse(tr.rand(78),Tensor[A]);# <- binds A=78 within this contextAs indicated above, the type-checker sees the shaped-tensor/array types. Additionally, these are subclasses of their rightful parents, so we can pass these to functions typed with vanillatorch.Tensorandnumpy.ndarryannotations, and type checkers will be a-ok with that.defvanilla_numpy(x:np.ndarray):...defvanilla_torch(x:tr.Tensor):...vanilla_numpy(arr)# type checker: OKvanilla_torch(arr)# type checker: Error!vanilla_torch(t1)# type checker: OKBasic forms of runtime validation performed byparse# runtime type checking>>>parse(1,Tensor[A])---------------------------------------------------------------------------ParseError:Expected<class'torch.Tensor'>, got: <class 'int'># dimensionality mismatch>>>parse(tr.ones(3),Tensor[A,A,A])---------------------------------------------------------------------------ParseError:shape-(3,)doesn't match shape-type (A=?, A=?, A=?)# unsatisfied shape pattern>>>parse(tr.ones(1,2),Tensor[A,A])---------------------------------------------------------------------------ParseError:shape-(1,2)doesn't match shape-type (A=1, A=1)# inconsistent dimension sizes across tensors>>>x,y=parse(...(tr.ones(1,2),Tensor[A,B]),...(tr.ones(4,1),Tensor[B,A]),...)---------------------------------------------------------------------------ParseError:shape-(4,1)doesn't match shape-type (B=2, A=1)To reiterate,parseis able to compare shapes across multiple tensors by entering into a "dimension-binding scope".
One can enter into this context explicitly:>>>fromphantom_tensorsimportdim_binding_scope>>>x=parse(np.zeros((2,)),NDArray[B])# binds B=2>>>y=parse(np.zeros((3,)),NDArray[B])# binds B=3>>>withdim_binding_scope:...x=parse(np.zeros((2,)),NDArray[B])# binds B=2...y=parse(np.zeros((3,)),NDArray[B])# raises!---------------------------------------------------------------------------ParseError:shape-(3,)doesn't match shape-type (B=2,)Support forLiteraldimensions:fromphantom_tensorsimportparsefromphantom_tensors.torchimportTensorimporttorchastrfromtyping_extensionsimportLiteralasLparse(tr.zeros(1,3),Tensor[L[1],L[3]])# static + runtime: OKparse(tr.zeros(2,3),Tensor[L[1],L[3]])#  # Runtime: ParseError - mismatch at dim 0Support forLiteraldimensions and variadic shapes:In Python 3.11 you can write shape types likeTensor[int, *Ts, int], where*Tsrepresents 0 or more optional entries between two required dimensions. phantom-tensor supports this "unpack" dimension. In this README we opt fortyping_extensions.Unpack[Ts]instead of*Tsfor the sake of backwards compatibility.fromphantom_tensorsimportparsefromphantom_tensors.torchimportTensorimporttorchastrfromtyping_extensionsimportUnpackasU,TypeVarTupleTs=TypeVarTuple("Ts")# U[Ts] represents an arbitrary number of entriesparse(tr.ones(1,3),Tensor[int,U[Ts],int)# static + runtime: OKparse(tr.ones(1,0,0,0,3),Tensor[int,U[Ts],int])# static + runtime: OKparse(tr.ones(1,),Tensor[int,U[Ts],int])# Runtime: Not enough dimensionsSupport forphantom types:Supports phatom type dimensions (i.e.intsubclasses that override__isinstance__checks):fromphantom_tensorsimportparsefromphantom_tensors.torchimportTensorimporttorchastrfromphantomimportPhantomclassEvenOnly(int,Phantom,predicate=lambdax:x%2==0):...parse(tr.ones(1,0),Tensor[int,EvenOnly])# static return type: Tensor[int, EvenOnly]parse(tr.ones(1,2),Tensor[int,EvenOnly])# static return type: Tensor[int, EvenOnly]parse(tr.ones(1,4),Tensor[int,EvenOnly])# static return type: Tensor[int, EvenOnly]parse(tr.ones(1,3),Tensor[int,EvenOnly])# runtime: ParseError (3 is not an even number)Compatibility with Runtime Type Checkersparseis not the only way to perform runtime validation using phantom tensors – they work out of the box with 3rd party runtime type checkers likebeartype! How is this possible?...We do something tricky here! At, runtimeTensor[A, B]actually returns aphantom type. This means thatisinstance(arr, NDArray[A, B])is, at runtime,actuallyperformingisinstance(arr, PhantomNDArrayAB), which dynamically generated and is able to perform the type and shape checks.Thanks to the ability bind dimensions within a specified context, allbeartypeneeds to do is faithfully callisinstance(...)within said context and we can have the inputs and ouputs of a phantom-tensor-annotated function get checked!fromtypingimportAnyfrombeartypeimportbeartype# type: ignoreimportpytestimporttorchastrfromphantom_tensors.alphabetimportA,B,Cfromphantom_tensors.torchimportTensorfromphantom_tensorsimportdim_binding_scope,parse# @dim_binding_scope:#   ensures A, B, C consistent across all input/output tensor shapes#   within scope of function@dim_binding_scope@beartype# <-- adds isinstance checks on inputs & outputsdefmatrix_multiply(x:Tensor[A,B],y:Tensor[B,C])->Tensor[A,C]:a,_=x.shape_,c=y.shapereturnparse(tr.rand(a,c),Tensor[A,C])@beartypedefneeds_vector(x:Tensor[Any]):...x,y=parse((tr.rand(3,4),Tensor[A,B]),(tr.rand(4,5),Tensor[B,C]),)z=matrix_multiply(x,y)z# type revealed: Tensor[A, C]withpytest.raises(Exception):# beartype raises error: input Tensor[A, C] doesn't match Tensor[A]needs_vector(z)# <- pyright also raises an error!withpytest.raises(Exception):# beartype raises error: inputs Tensor[A, B], Tensor[A, B] don't match signaturematrix_multiply(x,x)# <- pyright also raises an error!