v0.2.0

einx v0.2.0

Added

• Add partial support for Apple's mlx.
  • Supported:
    • einx.rearrange
    • einx.{elementwise|add|multiply|where|...}
    • einx.{reduce|sum|mean|...}
    • einx.{vmap_with_axis|flip|softmax|...}
  • Not supported yet:
    • einx.dot (mx.einsum is not implemented yet)
    • einx.vmap (mx.vmap does not fully support all primitives yet)
    • einx.{index|get_at|set_at|...} (due to relying on einx.vmap)
• Add partial support for dask.array.
  • Supported:
    • einx.rearrange
    • einx.{elementwise|add|multiply|where|...}
    • einx.{reduce|sum|mean|...}
    • einx.{vmap_with_axis|flip|softmax|...}
    • einx.dot
  • Not supported:
    • einx.vmap (vmap not implemented in dask)
    • einx.{index|get_at|set_at|...} (due to relying on einx.vmap)
• Add environment variable EINX_WARN_ON_RETRACE to warn when excessive retracing is detected.

Changed

• Allow -> and , to be composed with other operators. (This deprecates the existing [|] notation which should instead be implemented with
  composable ->. The feature is still maintained for backwards compatibility). For example:
  • einx.dot("b [c1->c2]", ...) expands to einx.dot("b [c1] -> b [c2]", ...)
  • einx.get_at("b p [i,->]", ...) expands to einx.get_at("b p [i], b p -> b p", ...)
• Allow einx.{set_at|add_at|...} to be called with zero-sized updates or coordinates (in which case the input tensor is returned as-is).
• Remove backend.dot which was not used anywhere but in the unit tests.
• Improve error reporting:
  • Drop internal stack frames when raising exceptions.
  • Better error when passing invalid shape constraints to einx functions.
• Reduce overhead of einx when using the PyTorch backend.

Fixed

• Fix compatibility of einx.nn.torch.Norm with PyTorch 2.2.
• Fix parameters in einn.param being ignored.
• Fix bug when using concatenations in einx.rearrange. See: #6
• Fix broadcasting new axes in einx.vmap_with_axis.