MLX updates

[altaic]

FYI, einsum just landed in mlx along with python array api support. Also, thanks for einx!

ml-explore/mlx#1269
ml-explore/mlx#1289

[fferflo]

Hi, thanks for bringing this up! Support for mx.einsum (einx.dot) is already implemented, although not in a release version yet: 0ad7af3 It will be included in the next release of einx. (Using some other einx functions such as einx.get_at with mlx is currently limited to the use cases that are supported by mx.vmap)