in hypothesis_gufunc/gufunc.py [0:0]
def gufunc_arg_shapes(signature, excluded=(), min_side=0, max_side=5, max_dims_extra=0):
"""Strategy to generate the shape of ndarrays for arguments to a function
consistent with its signature with extra dimensions to test broadcasting.
Parameters
----------
signature : str
Signature for shapes to be compatible with. Expects string in format
of numpy generalized universal function signature, e.g.,
`'(m,n),(n)->(m)'` for vectorized matrix-vector multiplication.
excluded : set(int)
Set-like of integers representing the positional for which the function
will not be vectorized. Uses same format as :obj:`numpy.vectorize`.
min_side : int or dict
Minimum size of any side of the arrays. It is good to test the corner
cases of 0 or 1 sized dimensions when applicable, but if not, a min
size can be supplied here. Minimums can be provided on a per-dimension
basis using a dict, e.g. ``min_side={'n': 2}``. One can use, e.g.,
``min_side={hypothesis_gufunc.gufunc.BCAST_DIM: 2}`` to limit the size
of the broadcasted dimensions.
max_side : int or dict
Maximum size of any side of the arrays. This can usually be kept small
and still find most corner cases in testing. Dictionaries can be
supplied as with `min_side`.
max_dims_extra : int
Maximum number of extra dimensions that can be appended on left of
arrays for broadcasting. This should be kept small as the memory used
grows exponentially with extra dimensions. By default, no extra
dimensions are added.
Returns
-------
shapes : list(tuple(int))
list of tuples where each tuple is the shape of an argument. Extra
dimensions for broadcasting will be present in the shapes.
Examples
--------
.. code-block:: pycon
>>> from hypothesis_gufunc.gufunc import BCAST_DIM
>>> gufunc_arg_shapes('(m,n),(n)->(m)',
min_side={'m': 1, 'n': 2}, max_side=3).example()
[(2, 3), (3,)]
>>> gufunc_arg_shapes('(m,n),(n)->(m)', max_side=9,
min_side={'m': 1, 'n': 2, BCAST_DIM: 5},
max_dims_extra=3).example()
[(6, 6, 7), (6, 7)]
>>> gufunc_arg_shapes('(m,n),(n)->(m)', excluded=(0,),
max_side=20, max_dims_extra=3).example()
[(11, 13), (1, 1, 1, 13)]
"""
_check_set_like(excluded, name="excluded")
min_side = _int_or_dict(min_side, 0)
max_side = _int_or_dict(max_side, DEFAULT_MAX_SIDE)
_order_check_min_max(min_side, max_side)
check_type(int, max_dims_extra, "extra dims")
order_check("extra dims", 0, max_dims_extra, GLOBAL_DIMS_MAX)
# Validate that the signature contains digits we can parse
weird_sig_digits = _weird_digits(signature)
if len(weird_sig_digits) > 0:
raise InvalidArgument("signature %s contains invalid digits: %s" % (signature, "".join(weird_sig_digits)))
# Parse out the signature: e.g., parses to [('n', 'm'), ('m', 'p')]
parsed_sig, _ = parse_gufunc_signature(signature)
# Get core shapes before broadcasted dimensions
shapes_st = _gufunc_arg_shapes(parsed_sig, min_side=min_side, max_side=max_side)
# Skip this broadcasting craziness if we don't want extra dims:
if max_dims_extra == 0:
return shapes_st
# We could use tuples instead without creating type ambiguity since
# max_dims_extra > 0 and avoid calling arrays, but prob ok like this.
bcast_dim_st = integers(min_value=min_side[BCAST_DIM], max_value=max_side[BCAST_DIM])
extra_dims_st = arrays(np.intp, (max_dims_extra,), elements=bcast_dim_st)
set_to_1_st = arrays(np.bool_, (len(parsed_sig), max_dims_extra))
# np.clip will convert to np int but we don't really care.
max_extra_per_arg = [
0 if nn in excluded else np.clip(GLOBAL_DIMS_MAX - len(ss), 0, max_dims_extra)
for nn, ss in enumerate(parsed_sig)
]
extra_per_arg_st = tuples(*[integers(min_value=0, max_value=mm) for mm in max_extra_per_arg])
return builds(_append_bcast_dims, shapes_st, extra_dims_st, set_to_1_st, extra_per_arg_st)