[Keras Ops] Add einops-style rearrange() to keras.ops

[DavidLandup0]

Context

Closes #20332.

Add an einops style rearrange() operation to the keras.src.ops.einops.
Making a new module for this makes sense if we plan to add other similar operations and want to re-use some of the private-marked utility methods here or generalize the logic further.

It makes use of solely reshape() and transpose(), so these natively run on all backends. Furthermore, only these two operations are applied so it should be fine with symbolic tensors (sometimes, einops doesn't work correctly with symbolic tensors in a computation graph).

Comparison with einops

Direct comparison with einops, using the examples from the official documentation:

import numpy as np
import keras
import einops

tensor = np.random.randn(2, 10, 10, 3)

examples = [
    (tensor, 'b h w c -> b h w c', {}),
    (tensor, 'b h w c -> b c h w', {}),
    (tensor, 'b h w c -> (b h) w c', {}),
    (tensor, 'b h w c -> h (b w) c', {}),
    (tensor, 'b h w c -> b (c h w)', {}),
    (tensor, 'b (h h1) (w w1) c -> b (h w1) (w h1) c', {'h1': 2, 'w1': 2}),
    (tensor, 'b (h h1) (w w1) c -> b c (h w1) (w h1)', {'h1': 2, 'w1': 2}),
    (tensor, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', {'h1': 2, 'w1': 2}),
    (tensor, 'b (h h1) (w w1) c -> b h w (c h1 w1)', {'h1': 2, 'w1': 2})
]

results = []

for data, pattern, axes_lengths in examples:
    keras_result = keras.ops.rearrange(data, pattern, **axes_lengths)
    einops_result = einops.rearrange(data, pattern, **axes_lengths)
    
    are_equal = np.allclose(keras_result, einops_result)
    
    result_entry = {
        "pattern": pattern,
        "from": data.shape,
        "to": keras_result.shape,
        "equal": are_equal
    }
    
    results.append(result_entry)

Results in:

{'pattern': 'b h w c -> b h w c', 'from': (2, 10, 10, 3), 'to': TensorShape([2, 10, 10, 3]), 'equal': True}
{'pattern': 'b h w c -> b c h w', 'from': (2, 10, 10, 3), 'to': TensorShape([2, 3, 10, 10]), 'equal': True}
{'pattern': 'b h w c -> (b h) w c', 'from': (2, 10, 10, 3), 'to': TensorShape([20, 10, 3]), 'equal': True}
{'pattern': 'b h w c -> h (b w) c', 'from': (2, 10, 10, 3), 'to': TensorShape([10, 20, 3]), 'equal': True}
{'pattern': 'b h w c -> b (c h w)', 'from': (2, 10, 10, 3), 'to': TensorShape([2, 300]), 'equal': True}
{'pattern': 'b (h h1) (w w1) c -> b (h w1) (w h1) c', 'from': (2, 10, 10, 3), 'to': TensorShape([2, 10, 10, 3]), 'equal': True}
{'pattern': 'b (h h1) (w w1) c -> b c (h w1) (w h1)', 'from': (2, 10, 10, 3), 'to': TensorShape([2, 3, 10, 10]), 'equal': True}
{'pattern': 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', 'from': (2, 10, 10, 3), 'to': TensorShape([8, 5, 5, 3]), 'equal': True}
{'pattern': 'b (h h1) (w w1) c -> b h w (c h1 w1)', 'from': (2, 10, 10, 3), 'to': TensorShape([2, 5, 5, 12]), 'equal': True}

Other

TODO:

• Write tests
• Performance benchmark

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 93.40659% with 6 lines in your changes missing coverage. Please review.

Project coverage is 81.98%. Comparing base (ab3c8f5) to head (90078cf).
Report is 7 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/ops/einops.py	94.38%	4 Missing and 1 partial ⚠️
keras/api/_tf_keras/keras/ops/__init__.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #20733      +/-   ##
==========================================
+ Coverage   81.95%   81.98%   +0.02%     
==========================================
  Files         553      554       +1     
  Lines       51446    51549     +103     
  Branches     7957     7972      +15     
==========================================
+ Hits        42164    42262      +98     
- Misses       7346     7347       +1     
- Partials     1936     1940       +4     

Flag	Coverage Δ
keras	81.80% <93.40%> (+0.02%)	⬆️
keras-jax	64.09% <92.30%> (+0.04%)	⬆️
keras-numpy	58.95% <92.30%> (+0.05%)	⬆️
keras-openvino	29.96% <92.30%> (+0.12%)	⬆️
keras-tensorflow	64.75% <93.40%> (+0.04%)	⬆️
keras-torch	64.13% <92.30%> (+0.03%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[fchollet]

Thanks for the PR!

[fchollet]

An op should be able to run on either symbolic Keras tensors or backend native eager tensors. And they should render as a single node in the op graph. This would require creating a class for the op, with a compute_output_spec method (see how other ops are implemented)

[DavidLandup0]

Ah, sorry, forgot to add the class.

[DavidLandup0]

@fchollet Thank you for the feedback! Changes should be reflected in the latest commit.
I'll add more detailed test cases and a performance benchmark still. Sorry for the delay on those.

[DavidLandup0]

Do we want any documentation or code comments on these?

[fchollet]

Thanks for the update!

[fchollet]

Please also add a code example.

[DavidLandup0]

Added examples to mirror: https://einops.rocks/api/rearrange/

[fchollet]

It's unusual to inline logic in a src/ops/ op rather than defining it N times in the backends in a backend specific fashion. But it's done for a couple other ops (image ops in particular). It's fine.

[DavidLandup0]

Yeah, I was debating opening N backend operations instead of one here. Though, since it just uses reshape() and transpose(), it gets to use backend-equal implementations by virtue of keras.ops by default. Figured that lower redundancy/copying is preferred in this case, especially since we could look into adding more operations in keras.src.ops.einops in the future.

[fchollet]

Please add this file to the list of excluded test files for openVINO, or otherwise fix the test. OpenVINO tests are failing.

[DavidLandup0]

Removed numpy in lieu of keras.ops - some of the ops in the tests themselves (all(), etc.) don't seem to be supported by openVINO. Skipped those.

[fchollet]

LGTM, thank you