def einsum(const_context, builder, op):
    '''
    MIL einsum is either
    - (B,C,H,W1) * (B,W1,H,W2) = (B,C,H,W2)
    or
    - (C,H,W1) * (W1,H,W2) = (C,H,W2)

    Hence to support it, first transpose the 2 inputs, so that the matrices
    to be multiplied are on the last 2 axes,
    then call bmm, and finally transpose the result again
    '''
    rank = op.values[0].rank
    perm = [0, 2, 1, 3] if rank == 4 else [1, 0, 2]
    input_names = make_input(const_context, builder, op.values)

    output_name_1 = op.name + "_transpose_1"
    output_name_2 = op.name + "_transpose_2"
    builder.add_transpose(name=op.name + "_transpose_x",
                          axes=perm,
                          input_name=input_names[0],
                          output_name=output_name_1
    )
    builder.add_transpose(name=op.name + "_transpose_y",
                          axes=perm,
                          input_name=input_names[1],
                          output_name=output_name_2
    )
    builder.add_batched_mat_mul(
        name=op.name + "_batch_matmul",
        input_names=[output_name_1, output_name_2],
        output_name=op.outputs[0].name + "_pre_transpose"
    )
    builder.add_transpose(name=op.name,
                          axes=perm,
                          input_name=op.outputs[0].name + "_pre_transpose",
                          output_name=op.outputs[0].name
    )

    MIL einsum is either
    - (B,C,H,W1) * (B,W1,H,W2) = (B,C,H,W2)
    or
    - (C,H,W1) * (W1,H,W2) = (C,H,W2)

    Hence to support it, first transpose the 2 inputs, so that the matrices
    to be multiplied are on the last 2 axes,
    then call bmm, and finally transpose the result again
    
coremltools/converters/mil/backend/nn/op_mapping.py