Repository: dask/dask
File: dask/array/tests/test_routines.py (Python)

Code before:
        (4,) * 2,
        (4,) * 2,
        (4,) * 2,
        (1,) * 8,
        (4,) * 2,
    )

    np_arr1 = np.random.random((2, 4, 4))
    np_arr2 = np.random.random((2, 4, 4))

    arr1 = da.from_array(np_arr1, chunks=(1, 2, 2))
    arr2 = da.from_array(np_arr2, chunks=(1, 2, 2))
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=da.PerformanceWarning)
        result = da.einsum("aij,amn->ijmn", arr1, arr2)
    assert result.chunks == ((1,) * 4,) * 4
    assert_eq(np.einsum("aij,amn->ijmn", np_arr1, np_arr2), result)

    # regression test for GH11627
    z = da.ones(
        shape=(40000, 2, 10, 2, 10), dtype=np.float64, chunksize=(40000, 1, 5, 2, 10)
    )
    x = da.ones(shape=(2, 10, 10), dtype=np.float64, chunksize=(2, 10, 10))
    y = da.ones(shape=(2, 10, 10), dtype=np.float64, chunksize=(2, 10, 10))
    res = da.einsum("abcde,bfc,dfe->acef", z, x, y)
    assert res.numblocks == 1


@pytest.mark.parametrize(
    "optimize_opts", [(True, False), ("greedy", False), ("optimal", False)]
)
def test_einsum_optimize(optimize_opts):
    sig = "ea,fb,abcd,gc,hd->efgh"
    input_sigs = sig.split("->")[0].split(",")
    np_inputs, da_inputs = _numpy_and_dask_inputs(input_sigs)

    opt1, opt2 = optimize_opts

    assert_eq(
        np.einsum(sig, *np_inputs, optimize=opt1),
        da.einsum(sig, *np_inputs, optimize=opt2),
    )

    assert_eq(
        np.einsum(sig, *np_inputs, optimize=opt2),
        da.einsum(sig, *np_inputs, optimize=opt1),
    )


@pytest.mark.parametrize("order", ["C", "F", "A", "K"])
def test_einsum_order(order):

Review comment (suggestion):
```suggestion
    assert res.numblocks == (1, 1, 1, 1)
```

Code after:
        (4,) * 2,
        (4,) * 2,
        (4,) * 2,
        (1,) * 8,
        (4,) * 2,
    )

    np_arr1 = np.random.random((2, 4, 4))
    np_arr2 = np.random.random((2, 4, 4))

    arr1 = da.from_array(np_arr1, chunks=(1, 2, 2))
    arr2 = da.from_array(np_arr2, chunks=(1, 2, 2))
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=da.PerformanceWarning)
        result = da.einsum("aij,amn->ijmn", arr1, arr2)
    assert result.chunks == ((1,) * 4,) * 4
    assert_eq(np.einsum("aij,amn->ijmn", np_arr1, np_arr2), result)

    # regression test for GH11627
    z = da.ones(
        shape=(40000, 2, 10, 2, 10), dtype=np.float64, chunksize=(40000, 1, 5, 2, 10)
    )
    x = da.ones(shape=(2, 10, 10), dtype=np.float64, chunksize=(2, 10, 10))
    y = da.ones(shape=(2, 10, 10), dtype=np.float64, chunksize=(2, 10, 10))
    res = da.einsum("abcde,bfc,dfe->acef", z, x, y)
    assert res.numblocks == (1, 1, 1, 1)


@pytest.mark.parametrize(
    "optimize_opts", [(True, False), ("greedy", False), ("optimal", False)]
)
def test_einsum_optimize(optimize_opts):
    sig = "ea,fb,abcd,gc,hd->efgh"
    input_sigs = sig.split("->")[0].split(",")
    np_inputs, da_inputs = _numpy_and_dask_inputs(input_sigs)

    opt1, opt2 = optimize_opts

    assert_eq(
        np.einsum(sig, *np_inputs, optimize=opt1),
        da.einsum(sig, *np_inputs, optimize=opt2),
    )

    assert_eq(
        np.einsum(sig, *np_inputs, optimize=opt2),
        da.einsum(sig, *np_inputs, optimize=opt1),
    )


@pytest.mark.parametrize("order", ["C", "F", "A", "K"])
def test_einsum_order(order):