Skip to content

Distributed einsum calls from python #463

@obackhouse

Description

@obackhouse

Are einsum calls from Python using some distributed layout not supported? The following code segfaults for me on >1 process

import numpy as np
import tiledarray as ta

world = ta.get_default_world()

ref = np.random.rand(4, 4)

x = ta.TArray([4, 4], 2, world=world)
for i, tile in enumerate(x):
    if i % world.size == world.rank:
        slices = tuple(slice(a, b) for a, b in zip(tile.range.start, tile.range.stop))
        tile.data = ref[slices].copy()
world.fence()

y = ta.TArray()
ta.einsum("ij,jk->ik", x, x, y)

as does something even simpler like ta.einsum("ij->ij", x, y).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions