Package: asta
+----+|asta|+----+Shape annotations for numpy arrays and pytorch/tensorflow tensors.Introduction------------This library defines subscriptable classes ``Array``, ``Tensor``, and``TFTensor`` for use in ``isinstance()`` checks and type annotations.  It alsoadds a decorator, ``@typechecked``, which implements toggleable static typeenforcement for the classes described above.Installation------------$ pip install asta$ pip install git+git://github.com/brendanxwhitaker/asta.gitBasics------Asta supports checking dtypes and/or shapes with wildcards and ellipses:>>> Array[float]                # dtype=np.float64.>>> Array[float, 1, 2, 3]       # dtype=np.float64, shape=(1,2,3).>>> Array[float, (1, 2, 3)]     # dtype=np.float64, shape=(1,2,3).>>> Array[1, 2, 3]              # shape=(1,2,3).>>> Array[int, ()]              # dtype=np.int64, shape=().>>> Array[str, 1, 2 ..., 3]     # dtype=unicode, shape=(1,2,*...*,3).>>> Array[str, 1, 2 -1, 3]      # dtype=unicode, shape=(1,2,*,3).>>> Tensor[torch.uint8, 1]      # dtype=torch.uint8, shape=(1,).>>> TFTensor[tf.complex128, ()] # dtype=tf.complex128, shape=().Wildcard dimensions (-1)------------------------An ``-1`` can be used as a wildcard in place of a single dimension size. Theybehave as a stand-in for any positive integer, and they can be used as manytimes as desired within a shape hint.Ellipses (...)--------------An ``Ellipsis`` object can be used in place of a nonnegative integer number ofpositive integer dimension sizes. The intended use case for this is when youknow the prefix of a tensor shape, the batch size for example, but not theremaining dimension sizes (perhaps they vary). You can use something like:>>> Array[32, ...]to match arrays with shape:>>> (32,)>>> (32, 24)>>> (32, 1, 2, 3)and so on. These should be a LAST RESORT, since as will be discussed below, wecan specify dimensions, shapes, or portions of shapes as variables, even ifthey will change at runtime, e.g. on every iteration of some loop.Multiple ellipses can be used within a single annotation as well, in case youhappen to know that a tensor or array will have a dimension size of 7 somewherein its shape, but you don't know where, or you don't know the total rank.Something like the following:>>> TFTensor[..., 7, ...]will match tensorflow tensors with shape:>>> (7,)>>> (1, 7)>>> (7, 3)>>> (1, 2, 3, 7, 4, 5, 6)and so on. Note that two ellipses cannot be used consecutively, e.g.>>> Array[1, 2, ..., ..., 5]since asta won't be able to determine which dimensions should be substitutedfor which ellipsis object.Typechecked decorator---------------------As mentioned above, asta implements a decorator for runtime typechecking usingtype hints like ``Array[]``. It can be used to decorate functions, instancemethods, class methods, metaclass methods, or even entire classes (this is justequivalent to decorating each of the class's methods).The following gives an example of using the ``@typechecked`` decorator toenforce torch tensor shapes and dtypes at runtime. The function ``kl`` willraise a TypeError if called with inputs which have any dtype other than``torch.float32``, or any shape other than ``(8, 64)``.>>> import os>>> import torch.nn.functional as F>>> from asta import Tensor, typechecked>>>>>> os.environ["ASTA_TYPECHECK"] = "1">>>>>>>>> @typechecked>>> def kl(t_1: Tensor[float, 8, 64], t_2: Tensor[float, 8, 64]) -> Tensor[()]:>>>     """ Computes the KL divergence of two Tensors of shape ``(8, 64)``. """>>>     divergence = F.kl_div(t_1, t_2, reduction="sum")>>>     return divergenceA runnable example is given in ``example.py`` in the repository root. For amore fleshed-out, real world example, see ``asta/tests/rl/`` for a typecheckedpolicy gradient implementation.Variable shapes and dimensions------------------------------Asta also supports variable shape arguments:>>> from asta import dims, shapes>>> BATCH_SIZE = dims.BATCH_SIZE>>> SEQ_LEN = dims.SEQ_LEN>>> OBS_SHAPE = shapes.OBS_SHAPE>>> ...>>> dims.BATCH_SIZE = 32>>> dims.SEQ_LEN = 256>>> shapes.OBS_SHAPE = (4, 5, 6)>>> ...>>> Tensor[float, BATCH_SIZE, SEQ_LEN]          # shape=(32,256).>>> Tensor[float, (BATCH_SIZE,) + OBS_SHAPE]    # shape=(32,4,5,6).These can be accessed from ``asta.dims`` for use in annotating top-levelfunctions and then set after import-time during the execution of a program,allowing for shape annotations which depend on control flow or which changeduring execution:(utils.py)>>> from asta import dims>>> HIDDEN_DIM = dims.HIDDEN_DIM>>> @typechecked>>> def identity(x: Tensor[float, HIDDEN_DIM]) -> Tensor[float, HIDDEN_DIM]:>>>     return x...(main.py)>>> from utils import example_fn>>> dims.HIDDEN_DIM = 768>>> x = torch.ones((768,))>>> y = identity(x)As seen above, this allows dimension or shape constants, stored in``asta.dims`` and ``asta.shapes`` respectively by setting a named attribute, topersist as globals across modules. This need not be constants, either. We canreset the value of ``dims.<NAME>``, and all typechecked function callsoccurring after the aforementioned statement executes will reflect the updatedvalue of ``dims.<NAME>``. Under the hood, this works because ``dims`` and``shapes`` are actually instances of classes with the ``__getattr__`` and``__setattr__`` functions overridden. They return a placeholder object for anyattribute accessed, and then update the value of that placeholder whenever thatattribute is set later on.Dimension inference-------------------Asta also supports arbitrary expressions of symbolic shape arguments:>>> from asta import symbols>>> X = symbols.X>>> Y = symbols.Y>>> Tensor[float, (X, X**2, X**3)]      # e.g. shape=(2,4,8) or shape=(1,1,1).>>> Tensor[float, (X + Y, (X + Y)**2)]  # Multivariate supported as well.Checks on the above objects will do inference based on the passed tensor,constructing a system of equations based on the actual shape and returning amatch if there exists at least one solution to the system. This means theequation solver is lenient, if you pass it a system which has many solutions,or infinitely many solutions, it will still return a match. Something like:>>> from asta import symbols>>> X = symbols.X>>> Y = symbols.Y>>> Tensor[X + Y]used as an annotation for ``torch.ones((50,))``, for example, yields theequation ``X + Y = 50``, which has many solutions for positive integer valuesof ``X`` and ``Y``. But this will not raise an error. Similar to the case withEllipses, symbols should only be used as a LAST RESORT, since it is oftenpossible to make a lazy-definition of a ``dims`` placeholder prior to thefunction call instead, which will be more precise. They are useful, however, ifyou're averse to peppering your code with dimension storage statements (like``dims.W = 256`` or something), or if it doesn't matter what the specificvalues of the sizes are, only that they obey some constraint relative to eachother.Configuration-------------The behavior of the ``@typechecked`` decorator can be configured through aconfiguration file, like ``setup.cfg``, ``pyproject.toml``, ``astarc``, or``.astarc``. The ``on`` option can also be toggled via an environment variablecalled ``ASTA_TYPECHECK``, which, if set to "1", will check, and will donothing if set to "0". If the option from the environment variable and theconfiguration file conflict, asta will default to ``off``.An example config file is given below, and in ``asta/defaults/astarc``.>>> [MASTER]>>>>>> on=yes>>> raise-errors=yes>>> print-passes=yes>>> check-non-asta-types=no>>> check-all-sequence-elements=yesAnd explanations of the options:``on`` : Determines whether or not functions are typechecked. Does as little aspossible other than calling the wrapped function when set to ``no``.``raise-errors`` : If ``yes`` errors will be raised when a typecheck fails,otherwise, the error will only be printed.``print-passes`` : If ``yes`` all passed typechecks will be printed, otherwise,they will be silent (similar to mypy/pylint output).``check-non-asta-types`` : If ``yes`` asta will attempt to check the types ofeverything passed to or returned from a function. A wide variety of typesare supported. Note: some of the helper functions for these have beenadapted from those in the package ``typeguard``. Otherwise, it will onlycheck ``Array``, ``Tensor``, and ``TFTensor``, and composites of these,like ``Dict[str, Tensor[1,2,3]]``.``check-all-sequence-elements`` : If ``yes``, it will check the types of allelements in iterable types like ``List[*]``. Otherwise, it will only check thefirst element in an attempt to be faster.Subscript arguments-------------------The valid subscript arguments for ``Array`` and ``Tensor`` are as follows:Types-----Array-----1. Any python type from the following list:a. intb. floatc. boold. complexe. bytesf. strg. datetimeh. timedelta2. Any numpy dtype, e.g. ``np.int64``.3. Omitted (no argument passed).Tensor------1. Any python type from the following list:a. intb. floatc. boold. bytes2. Any ``torch.Tensor``-supported torch dtype, e.g. ``torch.int64``.3. Omitted (no argument passed).TFTensor--------1. Any python type from the following list:a. intb. floatc. boold. bytes2. Any ``torch.Tensor``-supported tensorflow dtype, e.g. ``tf.int64``.3. Omitted (no argument passed).Shapes------1. Nonnegative integers.2. ``-1``, a wildcard for any positive integer size.3. Ellipses (``...``), a placeholder for any contiguous sequence ofpositive integer sizes, including zero-length sequences.4. ``()`` or ``Scalar``, which both indicate a scalar array or tensor.These are interchangeable.5. Dimensions from ``asta.dims``.6. Shapes from ``asta.shapes``.7. Symbols from ``asta.symbols``.8. Omitted (no argument passed).Shape constraints and best practices------------------------------------There is a key difference between the way scalar values are handled in numpyand the way they are handled in torch. Consider an array/tensor of shape``(2,2,2)``. When indexing the first element along each dimension, which shouldbe scalar, we call ``a[0,0,0]``, where ``a`` is our array. We do the same to atensor ``t``, and assign the result to variables ``x`` and ``y``, respectively:>>> a = np.zeros((2,2,2))>>> t = torch.zeros((2,2,2))>>> x = a[0,0,0]>>> y = t[0,0,0]What are the types of ``x`` and ``y``?>>> type(x)<class 'numpy.float64'>>>> type(y)<class 'torch.Tensor'>Interestingly enough, while ``a`` is of type ``np.ndarray``, ``x`` is of type``np.float64``, a subclass of float, while both ``t`` and ``y`` are tensors.Note that ``x`` is not an array:>>> isinstance(x, np.ndarray)False>>> isinstance(x, float)TrueAnd ``y`` is not a float:>>> isinstance(y, torch.Tensor)True>>> isinstance(y, float)FalseAsta does not attempt to rectify this discrepancy, and so the behavior of``Array`` and ``Tensor`` when it comes to scalars is slightly different. In theabove, ``x`` is not an instance of ``Array``, while ``y`` is an instance of``Tensor``, even though they are indexed in an identical manner.Use of the ellipsis placeholder (``...``) in asta is meant to mirror its usagein numpy/torch indexing. This is why we allow ``...`` to take the place of anemtpy shape. Note that indexing scalar arrays with an ellipsis returns thearray unchanged:>>> np.zeros(())array(0.)>>> a = np.zeros(())>>> a[...]array(0.)And adding an Ellipsis anywhere in an already-filled index tuple will return ascalar array with the expected value:>>> a = np.zeros((2,2,2))>>> a[0,0,0]0.0>>> a[0,0,0,...]array(0.)>>> a[0,0,...,0]array(0.)>>> a[...,0,0,0]array(0.)In contrast, we take the ``-1`` wildcard argument to represent only a singlepositive integer shape element. So if you wanted to all arrays with shape``(1,*,1)``, where ``*`` is nonempty, i.e. don't match arrays of shape``(1,1)``, but do match any of the following:1. ``(1, 1, 1)``2. ``(1, 2, 1)``3. ``(1, 2, 3, 4, 5, 1)``You would use ``Array[1,-1,...,1]``.Performance-----------The runtime checking functionality of asta is NOT meant to be used insituations where performance/speed is critical. Furthermore, use of the valuesof type hints within python code, which ``@typechecked`` decorator relies on,is not recommended; the ability to type hint in python is meant to be justthat, a hint. The usefulness of using the decorator is as a debugging ortesting step when working on large, complicated models or workflows. Many ofthe native numpy and pytorch array/tensor functions allow arbitrary shapedinputs, and it is easy for a malformed shape to pass unnoticed, with no effectsother than poor downstream performance or results. Asta is meant to be a crudeaide in dealing with this common problem, but by no means a comprehensive one.This having been said, the ``isinstance()`` checks used are relatively cheap,and shouldn't cause a serious slowdown outside of exceptional cases.The recommended usage of this library would be to annotate all criticalfunctions which take or return ndarrays/tensors, and decorate them with``@typechecked``. One could then add a CI test which sets the``ASTA_TYPECHECK`` environment variable to ``1`` and runs a sample workflow.Any incorrect dtypes or malformed shapes will raise a TypeError, and typecheckswhich pass will print to stdout. This behavior is intentional, and meant tohelp researchers avoid silent performance degradation due to leaving theenvironment variable set, which will cause a slight slowdown which wouldotherwise occur silently.========INTERNAL========Todo----- Add ``# type: ignore`` comments in test files. The ``[type-arg]`` andellipses errors will be ignored when the package is installed. They just needto be silenced within the package itself. (DONE)- Delete ``demo.py``. (DONE)- Implement ``-1`` wildcard shape element suppport. (DONE)- Add tests for ``Tensor``. (DONE)- Write examples in README. (DONE)- Add tests for empty arrays and tensors. (DONE)- Consider making Ellipses only match empty shapes and positive integer shapes,but not zero-valued shapes. This would be done under the assumption that mostpeople are not interested in working with empty arrays/tensors. And if theyare, they can use the ``.size`` attribute for an easy check. (DONE)- Add reprs. (DONE)- Fix tensor strategy. (DONE)- Add an option to disable typechecked decorator (default=disabled, ``.astarc``file).- Add environment variable for typechecking. (DONE)- Add tests for ``@typechecked``. (DONE)- Consider changing name of decorator to ``@shapechecked``. (NO)- Consider dropping the ``Scalar`` object. The less unfamiliar objects, thebetter. (NO)- Add more descriptive error if you pass torch dtype for an Array or numpydtype for a Tensor. (DONE)- Add ``CudaTensor`` class. (DONE)- Consider adding support for arbitrary shape/type constraints as subscriptarguments. They would be of the form:>>> def fn(shape: Tuple[int, ...]) -> bool:>>>     n: int>>>     for i, elem in enumerate(shape):>>>         if i == 0:>>>             n = elem>>>         elif elem == 2 * n:>>>             n = elem>>>         else:>>>             return False>>>     return TrueThe above constraint would only pass for shapes ``(n, 2n, 4n, 8n...)``.To enforce it, you would use ``Array[fn]``. (NO)- Consider allowing instantiation to support kwargs. (DETAILED BELOW)- Consider adding a delimiter in the typecheck decorator output to distinguishthe typecheck passes between different functions and signatures. Somethinglike ``===============<asta.typechecked_function()>===============``. (DONE)- Consider adding support for tensorflow. (DONE)- Consider adding support for passing ``()`` as in ``Array[int, ()]`` to denotescalar arrays instead of ``Scalar`` or ``None``. (DONE)- Write tests for ``Scalar``.- Write analogues of new ``Array`` tests for ``Tensor``. (DONE)- Consider reserving ``None`` shape and type for unintialized shape globals.(NO)Attempting to typecheck with them will raise an error warning that they areuninitialized. Then you can set your dim variable defaults in config moduleall to ``None``. (DONE)- Add uninitialized dimension size error to decorator. See preceding note.(DONE)- Add section in README about using ``asta.dims`` for programmatically-setdimension sizes in shape annotations.- Fix base class of ``_ArrayMeta`` and ``_TensorMeta`` so that type ignore isnot needed in ``decorators.py``.- Consider removing library-specific metaclasses.- Consider making ``parse_subscript`` a classmethod of ``SubscriptableMeta``.(NO)- Remove torch, tensorflow from requirements. You should be able to ues``Array`` without torch installed. (DONE)- When you try to import Tensor when torch is not installed, no error should beraised. But when you then try to use/subscript Tensor, an error should beraised. So it should still be a subscriptable type, but the subscriptionmethod should try to import torch. (DONE)- Consider adding a ``soft`` mode where the decorator only prints errorsinstead of raising them, so that a user could see and correct all of them atonce. (DONE)- Consider adding support for checking arbitrary attributes on arrays andtensors. For example, if working on a reinforcement learning problem, and youwanted to make sure all tensors have the same timestep index, you could seta``<torch.Tensor>.timestep: int`` attribute and then assert that they allmatch at function call time. (DONE)- This could have syntax like:>>> def product(>>>     t1: Tensor(float, 1,2,3, timestep=K),>>>     t2: Tensor(float, 1,2,3, timestep=K + 1),>>> ) -> Tensor(float, 1,2,3):This would require implementing an ``__init__()`` for array classes, whichwould just return an instance of the metaclass with the appropriate classvariables set. Alternatively, we could use a dictionary:>>> def product(>>>     t1: Tensor[float, 1,2,3, {"timestep": K}],>>>     t2: Tensor[float, 1,2,3, {"timestep": K + 1}],>>> ) -> Tensor[float, 1,2,3]:This is much better because it doesn't break the convention of usingsubscriptable objects for annotations. (DONE)- This would also require an implementation of non-contant shapes/variables.For example, if you know all your tensors will have length ``k``, but ``k``is variable, then you import ``K`` from some asta namespace. This will returna placeholder-like object with no set value. At typecheck time, it will beinferred from the first argument checked which makes use of this variable.The typechecker will assert that all other instances of ``K`` have the samevalue as this initialized one, and if soft checking is enabled, it will printout all of them along with an error message. (DONE)- Rewrite README.- Fix source code headers. (DONE)_ Add support for generics:- List (DONE)_ Sequence (DONE)- Dict (DONE)- Set (DONE)- Callable (DONE)- Write tests for decorated classes, attributes, equation solver, placeholders,instance/class/metaclass methods, and the above generic subscriptable types.- Resolve the discrepancy between default dtypes of tensors (float32 in TF),and default dtype given the primitive float argument (float64). What is themost natural behavior? Follow PLA. (DONE)- In line with above, consider making checks for primitive types (int, float,complex, bytes, str) less strict. This should be done if there is time. (NO)- Put on PyPI.- Bug: disabling/enabling asta per-file is likely broken, since importing checksets the value of an environment variable, which is global during pythonexecution. (DONE)- Consider adding ``asta.shp`` and ``asta.dims`` so that integer and tupledimensions/shapes are in two separate storage modules. This would also meanyou could write ``shp.OBS`` intead of ``OBS_SHAPE`` or ``dims.OBS_SHAPE``,which are both rather long and take up a good chunk of the hint size.- Decide on whether the storage modules should be plural (shp, dim, symbol) or(shps, dims, symbols)? Singular is shorter and looks nicer, maybe?- Placeholder should support addition as tuple concatenation. Does this alreadywork?- Placeholders dims (scalars) should support arbitrary mathematical expressionsinvolving symbols.- Add configuration parsing. (DONE)- Consider just not implementing per-module typecheck enable/disable. It shouldbe all or nothing via either the configuration file option or the environmentvariable.- Consider implementing per-module disable via a global variable in eachmodule. The decorators can see the scope of the module from which they arecalled, and these scopes are independent, so set ``ASTA_MODULE_DISABLE =True``and asta will skip checking that module.- Don't support default precisions because it decreases type transparency.(DONE)- Support subcriptable aliases:>>> CudaTensor = Tensor[{"device": "cuda:0"}]>>> FloatCudaTensor = CudaTensor[float]Currently, subscripting ``CudaTensor`` as above will overwrite/delete theexisting subscript. So just add a class attribute that saves the subscript,and upon subscription, combines the previous subscript with the newsubscript, raising an error if an attribute is already set. Actually, onsecond though, instead of saving the subscript, just save the parsed classattributes. (DONE)- Bug: tuples cannot be summed without unpacking ``Tensor[shapes.OB + (1,)]``.(DONE)- Bug: Header should still be printed when an error is raised if``ox.print_passes`` is off. (DONE)- ---MERGE---(DONE)- Update README (DONE)- Add warning when ``ox.print_passes`` is False, so that the user doesn'taccidentally leave it on, slowing down their program, and make this happen atimport-time, so that if the part where they're calling typechecked functionsdoesn't happen for a while, they still know immediately. (DONE)- ---MERGE---- Consider adding an ``args`` module. So if you have a function:>>> def fn(k: int) -> Tensor[<k>]:>>>     return torch.ones((k,))where the shape is dependent on the arguments, you could annotate it.Syntax would look like this:>>> def fn(k: int) -> Tensor[args.k]:>>>     return torch.ones((k,))- Consider adding probabilistic checking, i.e. if a function has been called100,000 times and none of them have failed, maybe stop checking it,especially if it's taking a long time. We can do inference and estimate theprobability that the function will fail.- Before we do the above, attempt to jit compile everything in asta.- Consider making overwrites of class attributes like:>>> Tensor[float][float]raise an error saying the attribute has already been defined. Perhaps``kwattrs`` allows overwriting as long as there are no keys in common orsomething. Is an error expected behavior, or should we be lenient?- Consider adding checks before looping over any iterable to make sure thesubscript annotation is an astatype in the case where``ox.check_non_asta_types`` is false.- Consider making ``@typechecked`` callable so we can call:>>> @typechecked(print=True)>>> def ...So then we can have ``print-all-passes=no`` but still print the one we'redebugging.Acknowledgements----------------- Based on the excellent packages 'nptyping' by Ramon Hagenaars and 'typeguard'by Alex Grönholm.- Thanks to Michael Crawshaw (@mtcrawshaw) for helpful comments on naturalshape constraints and handling of 0-valued dimension sizes.