Package: shapecheck
ShapeCheckFramework-agnostic library for checking array/tensor shapes at runtime.Finding the root of shape mismatches can be troublesome, especially with
broadcasting rules and mutable arrays. Comments documenting shapes can easily
become out of date as code evolves. This library aims to solve both of those
problems by ensuring function input/output shape expectations are met. The
concise syntax for expressing shapes serves to document code as well, so new
users can quickly understand what's going on.With frameworks like JAX or TensorFlow, "runtime" is actually "compile" or
"trace" time, so you don't pay any cost during execution. For frameworks like
PyTorch, asynchronous execution will hide the cost of shape checking. You only
pay a small overhead with synchronous, eager frameworks like numpy.Install LibraryFrom PyPI:pipinstall--upgradeshapecheckTo install the latest version:pipinstall--upgradegit+https://github.com/n2cholas/shapecheck.gitUsageimportnumpyasnpfromshapecheckimportcheck_shapes@check_shapes({'imgs':'N,W,W,-1','labels':'N,1'},aux_info=None,out_='')defloss_fn(batch,aux_info):# do something with aux_infodiff=(batch['imgs'].mean((1,2,3))-batch['labels'].squeeze())returnnp.mean(diff**2)loss_fn({'imgs':np.ones((3,2,2,1)),'labels':np.ones((3,1))},np.ones(1))loss_fn({'imgs':np.ones((5,3,3,4)),'labels':np.ones((5,1))},'any')# Below line fails:loss_fn({'imgs':np.ones((3,5,2,1)),'labels':np.ones((3,1))},'any')Error message:shapecheck.exception.ShapeError: in function loss_fn.
Named Dimensions: {'N': 3, 'W': 5}.
Input:
    Argument: batch  Type: <class 'dict'>
        MisMatch: Key: imgs Expected Shape: ('N', 'W', 'W', -1) Actual Shape: (3, 5, 2, 1).
        Match:    Key: labels Expected Shape: ('N', 1) Actual Shape: (3, 1).
    Skipped:  Argument: aux_info.In the above example, we compute the loss with a batch of data, which is a
dictionary with images and labels. We specify that we wantNsquare images
which can have any number of channels (indicated by the-1).  Inputs tocheck_shapecan be arbitrarily nested dicts/lists/tuples, as long as the
structure of the shape specification matches the structure of the inputs to the
decorated function.We also specify thataux_infoshouldn't be checked. Equivalently, we could've
excluded it from the definition:@check_shapes({'imgs':'N,W,W,-1','labels':'N,1'},out_='')or passed it as a positional argument.@check_shapes({'imgs':'N,W,W,-1','labels':'N,1'},None,out_='')Finally, we specify the output shape should be a scalar viaout_=''. All
non-input shape arguments tocheck_shapehave an underscore after them so
they don't conflict with the decorated function's arguments (for now, justout_andmatch_callees_).If you have a function with shape-checking that calls many other functions with
shape-checking, you can optionally enforce that dimensions with the same letter
name in the caller correspond to the same sized dimension in the callees viamatch_callees_=True.  That is, you can check that a function's input named
dimensions match the same named dimensions of all checked functions higher in
the call stack. For example:@check_shapes('M','N','O',out_='M')defcallee(a,b,c):returna@check_shapes('M','N','R')defcaller_fn_1(x,y,z):returncallee(y,x,z)@check_shapes('M','N','R',match_callees_=True)defcaller_fn_2(x,y,z):returncallee(y,x,z)caller_fn_1(np.ones(5),np.ones(6),np.ones(7))# succeedscaller_fn_2(np.ones(5),np.ones(6),np.ones(7))# failsHere, we (accidentally) swappedxandywhen callingcallee.caller_fn_1succeeds because the inputs are compatible when considering the
named dimensions forcallee_fnalone. Butcaller_fn_2fails because the
named dimensions are inconsistent between the caller and the callee. The
following error would be produced:shapecheck.exception.ShapeError: in function callee.
Named Dimensions: {'M': 5, 'N': 6, 'R': 7, 'O': 7}.
Input:
    MisMatch: Argument: a Expected Shape: ('M',) Actual Shape: (6,).
    MisMatch: Argument: b Expected Shape: ('N',) Actual Shape: (5,).
    Match:    Argument: c Expected Shape: ('O',) Actual Shape: (7,).This library also supports variadic dimensions. You can use '...' to indicate 0
or more dimensions:@check_shapes('dim, ..., 1','..., dim, 1')defg(a,b):passg(np.ones((2,3,4,1)),np.ones((5,2,1)))# succeedsg(np.ones((3,1)),np.ones((3,1)))# succeedsg(np.ones((2,3,4,1)),np.ones((1,1)))# failsThe last statement fails with the following error, sincedimdoesn't match:shapecheck.exception.ShapeError: in function g.
Named Dimensions: {'dim': 2}.
Input:
    Match:    Argument: a Expected Shape: ('dim', '...', 1) Actual Shape: (2, 3, 4, 1).
    MisMatch: Argument: b Expected Shape: ('...', 'dim', 1) Actual Shape: (1, 1).You can also name the variadic dimensions, to ensure that a contiguous sequence
of dimensions match between arguments. For example:@check_shapes('batch,variadic...','variadic...')defh(a,b):passh(np.ones((7,1,2)),np.ones((1,2)))# succeedsh(np.ones((6,2)),np.ones((1,1)))# failsh(np.ones((6,2)),np.ones((1)))# failsYou can enable/disable shapechecking globally as shown below:fromshapecheckimportis_checking_enabled,set_checking_enabledassertis_checking_enabled()set_checking_enabled(False)assertnotis_checking_enabled()set_checking_enabled(True)assertis_checking_enabled()Or via a context manager:assertis_checking_enabled()withset_checking_enabled(False):assertnotis_checking_enabled()assertis_checking_enabled()If you have any questions or issues with the library, please raise an issue onGitHub. Hope you enjoy using
the library!