in hypothesis_gufunc/gufunc.py [0:0]
def _append_bcast_dims(core_dims, b_dims, set_to_1, n_extra_per_arg):
"""Add extra broadcast dimensions to core dimensions of array shapes.
Parameters
----------
core_dims : list of tuples of int
list of tuples where each tuple is the core shape of an argument. It
has length `n_args`.
b_dims : ndarray of shape (max_dims_extra,)
Must be of `int` dtype and >= 0. Extra dimensions to pre-pend for
roadcasting.
set_to_1 : ndarray of shape (n_args, max_dims_extra)
Must be of `bool` dtype. Which extra dimensions get set to 1 for
broadcasting.
n_extra_per_arg : like-like of shape (n_args,)
Elements must be of int type. Must be in [0, max_dims_extra], how many
extra dimensions to pre-pend to each argument.
Returns
-------
shapes : list of tuples of int
list of tuples where each tuple is the shape of an argument. Extra
dimensions for broadcasting will be present in the shapes. It has
length `n_args`.
"""
# Build 2D array with extra dimensions
# e.g., extra_dims = [[2 5], [2 5]]
extra_dims = np.tile(b_dims, (len(core_dims), 1))
# e.g., extra_dims = [[1 5], [2 5]]
extra_dims[set_to_1] = 1 # This may be outside [min_side, max_side]
# Get full dimensions (core+extra), will chop some on left randomly
# e.g., shapes = [(5, 1, 3), (2, 5, 3, 1)]
# We use pp[len(pp) - nn:] instead of pp[-nn:] since that doesn't handle
# corner case with nn=0 correctly (seems like an oversight of py slicing).
# Call tolist() before tuple to ensure native int type.
shapes = [tuple(pp[len(pp) - nn :].tolist()) + ss for ss, pp, nn in zip(core_dims, extra_dims, n_extra_per_arg)]
return shapes