EinsumLayer

Layer that computes an einsum and optional bias.

Abstract Signature:

EinsumLayer(equation: str)

PyTorch

API: —
Strategy: Custom / Partial

Keras

API: keras.layers.EinsumDense
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: flax.nnx.Einsum
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: praxis.layers.Einsum
Strategy: Direct Mapping
Official Docs ↗