Repository: apple/coremltools

Issue #1537: Add support for Einsum equation format: hbg,heg->hbe
Created: 2022-06-18T00:54:18Z
🌱 Describe your Feature Request
I'm trying to convert vqgan-jax to Core ML format. I've successfully converted it to a TensorFlow model, but when I try to convert the TF model to Core ML I get the following error:
----> 3 mlmodel = ct.convert('/content/vqgan-saved-model', convert_to="mlprogram", source="tensorflow", compute_precision=ct.precision.FLOAT32)

11 frames
[/usr/local/lib/python3.7/dist-packages/coremltools/converters/mil/frontend/_utils.py](https://amz1rkk4g9d-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20220615-060045-RC00_455067423#) in build_einsum_mil(a_var, b_var, equation, name)
    104     else:
    105         raise NotImplementedError(
--> 106             "Einsum unsupported equation format: ", equation
    107         )
    108 

NotImplementedError: ('Einsum unsupported equation format: ', 'hbg,heg->hbe')

How can this feature be used?
Via the existing coremltools.convert(...) API.
Describe alternatives you've considered
N/A
Additional context
N/A

Duplicate issue #1251: Cannot convert keras.layers.MultiHeadAttention
Created: 2021-07-10T02:44:30Z
🐞Describe the bug
There is an issue when converting TF2 Keras models that contain MultiHeadAttention:
layers.MultiHeadAttention(num_heads=num_heads, key_dim=head_dim, name="attention")
The conversion fails with: ValueError: Cannot add const [512*is10, 512]
The is10 variable increments each time I try.
The problem seems to be when calculating the matrix size for one of the Einsums. I can't tell if it's the Q,K,V einsums causing trouble or the scaled dot product einsums.
I tried also using the MultiHeadAttention from TensorFlow addons, but that one failed with unsupported einsums.
The model trains and executes fine so this seems to be a conversion issue. I tried 5.0b2.
Trace
coremltools/converters/mil/frontend/tensorflow/ops.py in Einsum(context, node)
    493     a = context[node.inputs[0]]
    494     b = context[node.inputs[1]]
--> 495     x = build_einsum_mil(a, b, equation, node.name)
    496     context.add(node.name, x)
    497 

coremltools/converters/mil/frontend/_utils.py in build_einsum_mil(a_var, b_var, equation, name)
     66         if parsed_vectors_rev == ([0,1,2],[2,3,4],[0,1,3,4]):
     67              a_var, b_var = _swap(a_var, b_var)
---> 68         x_1 = mb.reshape(x=a_var, shape=[a_var.shape[0] * a_var.shape[1], a_var.shape[2]])
     69         x_2 = mb.reshape(x=b_var, shape=[b_var.shape[0], b_var.shape[1] * b_var.shape[2]])
     70         x = mb.matmul(x=x_1, y=x_2, transpose_x=False, transpose_y=False)

coremltools/converters/mil/mil/ops/registry.py in add_op(cls, **kwargs)
     59             @classmethod
     60             def add_op(cls, **kwargs):
---> 61                 return cls._add_op(op_cls, **kwargs)
     62 
     63             setattr(Builder, op_type, add_op)

coremltools/converters/mil/mil/builder.py in _add_op(cls, op_cls, **kwargs)
    160         # Shallow copy list inputs to ensure op inputs are immutable
    161         kwargs = {k: v if not isinstance(v, (list, tuple)) else v[:] for k, v in kwargs.items() if v is not None}
--> 162         kwargs.update(cls._create_vars(
    163             input_spec=op_cls.input_spec,
    164             op_name=kwargs["name"], before_op=before_op,

coremltools/converters/mil/mil/builder.py in _create_vars(cls, input_spec, op_name, before_op, candidate_kv)
    143             if isinstance(in_type, (ScalarOrTensorInputType,
    144               ListOrScalarOrTensorInputType)):
--> 145                 var = cls._add_const(val, new_var_name, before_op)
    146                 update_dict[k] = var
    147 

~/.virtualenvs/codepredictor/lib/python3.9/site-packages/coremltools/converters/mil/mil/builder.py in _add_const(cls, val, name, before_op)
     73     def _add_const(cls, val, name, before_op):
     74         if not is_python_value(val):
---> 75             raise ValueError("Cannot add const {}".format(val))
     76         if any_symbolic(val):
     77             msg = (

ValueError: Cannot add const [512*is10, 512]
To Reproduce
The model source here reproduces this bug: https://github.com/keras-team/keras-io/blob/master/examples/generative/text_generation_with_miniature_gpt.py
System environment:

coremltools version: 5.0b2
OS: MacOS
macOS version: 11.4
XCode version: 12.5.1
How you install python: virtualenv
python version: 3.9.4
TensorFlow: 2.5

Closing comment (2022-07-14T19:21:23Z): Sorry @netanellavisdris - no updates.
This is related to #1537.