einx
einx copied to clipboard
fferflo
Universal Tensor Operations in Einstein-Inspired Notation for Python.
FYI, `einsum` just landed in `mlx` along with python array api support. Also, thanks for `einx`! https://github.com/ml-explore/mlx/pull/1269 https://github.com/ml-explore/mlx/pull/1289
When working with some types of data, I need to reshape arrays with zero elements. For example, ```python import numpy as np import einx import einops arr = np.zeros((0, 0))...
I am trying the provided examples, particularly those associated with neural networks. In particular, the following example does not seem to work: ```python x = einx.multiply("... [c]", x, einn.param(init=1e-5)) ```...
I am trying to get the example: ``` einx.add_at("b ([h w]) c, ([2] b) i, c i -> c [h w] b", image, coordinates, updates) ``` to work. Could you...
The `argmin` and `argmax` functions would be convenient among the set of `einx` functions. It seems obvious, but they should work just like `min` and `max` except that they would...
Current ways of squeezing an array on a given dimension using `rearrange` are : ``` vals_squeezed = einx.rearrange('a 1 c -> a c', vals) # dim 'k' is not named...
Attempting to use `einx.vmap` with an equinox module as the op argument will crash if the module has any parameters (as jax arrays are unhashable). This is particularly problematic, as...
The [einx.rearrange](https://github.com/fferflo/einx/blob/b7a90b413f74f7d215d2deedd25a760057e4ee2d/einx/op/rearrange.py#L92) function returns a `Union[einx.Tensor, Tuple[einx.Tensor, ...]]` and this is problematic when assigning the result into a variable with an `einx.Tensor` typehint, because the function might be returning a...
Resolves #15 This does introduce a dependency to cachetools to use an LRU cache with custom keys. Further, a dummy pass is still required through any module being vmapped over....
Hi there, I'm working on a feedstock PR so `einx` is installable as a conda package via conda-forge. @fferflo would you be open to being a maintainer? Usually that just...
