add evaluate method for einsum

[Mohamed-Ashraf273]

Details:

• Identified that the einsum operation was not functioning properly due to the absence of an overridden evaluate method and no optimized kernel.

• Implemented the evaluate method for einsum to ensure correct execution within loop nodes.

• Tested the fix.

Tickets:

GSoC2025 – pipeline enablement on Keras with OpenVINO backend
@rkazants

[mmikolajcz]

Hi @Mohamed-Ashraf273, could you share reproducer to this issue?

In OV Einsum and many other ops don't have evaluate in core implementation by design to reduce binary size.
Einsum inference should be handled by op decomposition transformation, however it may have failed for your case due to some bug/limitation that would be worth to investigate further

[Mohamed-Ashraf273]

Hi @Mohamed-Ashraf273, could you share reproducer to this issue?

In OV Einsum and many other ops don't have evaluate in core implementation by design to reduce binary size. Einsum inference should be handled by op decomposition transformation, however it may have failed for your case due to some bug/limitation that would be worth to investigate further

import numpy as np
from keras import ops
from openvino import Model, compile_model

index = ops.convert_to_tensor(0, np.int32)
arr = ops.convert_to_tensor(np.array([[1, 2], [3, 4]]))
mask = ops.convert_to_tensor(np.array([[1, 0], [0, 1]]))
loop_vars = (index, arr, mask)

index = ops.convert_to_tensor(0, np.int32)
arr = ops.convert_to_tensor(np.array([[1, 2], [3, 4]]).astype(np.float16))
mask = ops.convert_to_tensor(np.array([[1, 0], [0, 1]]).astype(np.float16))
# Add two tensors for einsum
a = ops.convert_to_tensor(np.random.randn(1, 1, 2048).astype(np.float16))
b = ops.convert_to_tensor(np.random.randn(1, 2048, 256).astype(np.float16))
loop_vars = (index, arr, mask, a, b)

def cond(index, arr, mask, a, b):
    row_idx = ops.mod(index, 2)
    arr_slice = ops.slice(arr, [row_idx, 0], [1, 2])
    sum_slice = ops.sum(arr_slice, axis=[0, 1])
    index_lt_4 = ops.less(index, 4)
    sum_lt_30 = ops.less(sum_slice, 30)
    return ops.logical_and(index_lt_4, sum_lt_30)


def body(index, arr, mask, a, b):
    row_idx = ops.mod(index, 2)
    mask_slice = ops.slice(mask, [row_idx, 0], [1, 2])
    arr_slice = ops.slice(arr, [row_idx, 0], [1, 2])
    mask_slice = ops.reshape(mask_slice, [1, 2])
    arr_slice = ops.reshape(arr_slice, [1, 2])
    arr_slice_updated = ops.where(mask_slice == 1, ops.multiply(arr_slice, 2), arr_slice + 5)
    arr_parts = []
    for i in range(2):
        current_row_idx = ops.mod(index, 2)
        is_target_row = ops.equal(i, current_row_idx)
        part = ops.where(is_target_row, arr_slice_updated, ops.slice(arr, [i, 0], [1, 2]))
        arr_parts.append(part)
    arr_updated = ops.concatenate(arr_parts, axis=0)
    mask_next = 1 - mask
    arr_final = arr_updated + index

    # Use einsum_out: add its sum to arr_final (broadcasted)
    einsum_out = ops.einsum('btd,ndh->btnh', a, b)
    einsum_sum = ops.sum(einsum_out, axis=[0, 1, 2, 3])  # scalar
    arr_final = arr_final + einsum_sum  # broadcast add

    return (index + 1, arr_final, mask_next, a, b)

result = ops.while_loop(cond, body, loop_vars, maximum_iterations=4)
final_index, final_arr, final_mask = result[:3]
ov_outputs = [final_index.output, final_arr.output, final_mask.output]

ov_model = Model(results=ov_outputs, parameters=[])
#save_model(ov_model, "while_loop_model.xml")
compiled = compile_model(ov_model, "CPU")
outputs = compiled({})

print("Final index:", outputs[0])
print("Final array:\n", outputs[1])
print("Final mask:\n", outputs[2])

(env) mohamed-ashraf@mohamed-ashraf-LOQ-15IRX9:~/Desktop/GSoC2025$ python test_while.py 
/home/mohamed-ashraf/Desktop/GSoC2025/env/lib/python3.12/site-packages/openvino/runtime/__init__.py:10: DeprecationWarning: The `openvino.runtime` module is deprecated and will be removed in the 2026.0 release. Please replace `openvino.runtime` with `openvino`.
  warnings.warn(
Traceback (most recent call last):
  File "/home/mohamed-ashraf/Desktop/GSoC2025/test_while.py", line 61, in <module>
    compiled = compile_model(ov_model, "CPU")
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mohamed-ashraf/Desktop/GSoC2025/env/lib/python3.12/site-packages/openvino/_ov_api.py", line 723, in compile_model
    return core.compile_model(model, device_name, {} if config is None else config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mohamed-ashraf/Desktop/GSoC2025/env/lib/python3.12/site-packages/openvino/_ov_api.py", line 623, in compile_model
    super().compile_model(model, device_name, {} if config is None else config),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Exception from src/inference/src/cpp/core.cpp:109:
Exception from src/inference/src/dev/plugin.cpp:53:
Exception from src/core/src/model.cpp:570:
Evaluation failed on opset7::Einsum Einsum_201 (opset1::Constant Constant_14[0]:f16[1,1,2048], opset1::Constant Constant_16[0]:f16[1,2048,256]) -> (f16[1,1,1,256])

Hi @mmikolajcz
This code raises that error when I try to use an einsum operation inside a while loop, but it works and gives correct results without the einsum operation.
I traced the error in openvnio codebase and found that it returns false due to the missing evaluate method for einsum.
The while_loop implementation is just an opset.loop node.
The same error occurred for gather_nd as well, due to the missing evaluate method.

[mmikolajcz]

Thanks @Mohamed-Ashraf273, I was able to reproduce it using keras master patched with https://github.com/keras-team/keras/pull/21369/files#diff-b15baa80baacd1d26bf238fc6c762baefee763b860076b8ee80308afd80cfa82
I've experimented with it a bit and it seem to be some kind of issue with common transformations not being applied when decomposed op is placed within loop body. Placing REGISTER_PASS(manager, EinsumDecomposition) before ConstantFolding in line:

openvino/src/common/transformations/src/transformations/common_optimizations/moc_transformations.cpp

Line 174 in c458a6b

REGISTER_PASS(manager, ConstantFolding)

seem to solve this problem. Einsum placed outside loop body seem to work correctly.
@itikhono from your experience with transformations, have you encountered something similar or do you know how to proceed?

[Mohamed-Ashraf273]

Thanks @Mohamed-Ashraf273, I was able to reproduce it using keras master patched with https://github.com/keras-team/keras/pull/21369/files#diff-b15baa80baacd1d26bf238fc6c762baefee763b860076b8ee80308afd80cfa82 I've experimented with it a bit and it seem to be some kind of issue with common transformations not being applied when decomposed op is placed within loop body. Placing REGISTER_PASS(manager, EinsumDecomposition) before ConstantFolding in line:

openvino/src/common/transformations/src/transformations/common_optimizations/moc_transformations.cpp

Line 174 in c458a6b

REGISTER_PASS(manager, ConstantFolding)

seem to solve this problem. Einsum placed outside loop body seem to work correctly.
@itikhono from your experience with transformations, have you encountered something similar or do you know how to proceed?

Thanks for the solution. It worked.