CODEWITHSUNDEEP
pythonnumpynumpy-einsum
fhchl
fhchl

Reputation: 751

Output shape of numpy.einsum

Is there a elegant way to precompute the shape of the result from np.einsum given einsum's input arguments (without running the computation)?

# Given a, b and signature with 
# a.shape == (1, 2, 5)
# b.shape == (4, 5)
einsum_shape('ijk,mk->ik', a, b) # returns (1, 5)

Upvotes: 1

Views: 221

Answers (2)

fhchl
fhchl

Reputation: 751

Based on @Divakar's answer, I came up with the following which is a bit more readable and raises errors, if unsupported subscript strings are passed.

def einsum_outshape(subscripts, *operants):
    """Compute the shape of output from `numpy.einsum`.

    Does not support ellipses.

    """
    if "." in subscripts:
        raise ValueError(f'Ellipses are not supported: {subscripts}')

    insubs, outsubs = subscripts.replace(",", "").split("->")
    if outsubs == "":
        return ()
    insubs = np.array(list(insubs))
    innumber = np.concatenate([op.shape for op in operants])
    outshape = []
    for o in outsubs:
        indices, = np.where(insubs == o)
        try:
            outshape.append(innumber[indices].max())
        except ValueError:
            raise ValueError(f'Invalid subscripts: {subscripts}')
    return tuple(outshape)

Upvotes: 0

Divakar
Divakar

Reputation: 221634

Here's something that works for generic number of inputs and relevant einsum-expression and also works for the specific scalar reduction case -

def einsum_outshape(einsum_expr, inputs):
    shps = np.concatenate([in_.shape for in_ in inputs])
    p = einsum_expr.split(',')
    s = p[:-1] + p[-1].split('->')
    if s[-1]=='':
        return ()
    else:
        inop = list(map(list,s))
        return tuple(shps[(np.concatenate(inop[:-1])[:,None]==inop[-1]).argmax(0)])

Sample run -

In [42]: a = np.random.rand(1,2,5)
    ...: b = np.random.rand(4,5)
    ...: c = np.random.rand(5,7,8)
    ...: d = np.random.rand(7,9)

In [43]: einsum_outshape('ijk,mk,kpq,pr->ikpqr', inputs=(a,b,c,d))
Out[43]: (1, 5, 7, 8, 9)

# Reduction to a scalar
In [44]: einsum_outshape('ijk,mk,kpq,pr->', inputs=(a,b,c,d))
Out[44]: ()

Upvotes: 1

Related Questions