reject clever
  • m_f@discuss.onlineEnglish
    2·
    1 month ago

    Great follow-up to “I don’t like NumPy”, it’s always better to show by example how things could be better. Minor nit, it would be nice if the different code samples at the end were displayed side-by-side somehow. Also wondering if the assert x.ndim == 1 in softmax_dumpy somehow meaningful? The other examples, including “DumPy (alternate)”, don’t have that check.

    • dynomight@lemmy.worldOPMEnglish
      3·
      1 month ago

      At one point, I actually had some (LLM-generated) boxes where you could click to switch between the different implementations for the same problem. But in the end I didn’t like how it looked, so I switched to simple expandy-boxes. Design is hard…

      There’s no magical significance to the assert x.ndim==1 check. I think I just wanted to demonstrate that the softmax code was “simple” and didn’t have to think about high dimensions. I think I’ll just remove that, thanks.

  • YetiBeets@lemmy.worldEnglish
    2·
    1 month ago

    Along the same vein, could we just move entirely to Einstein summation? It seems like your solution is 90% there.

    I assume there is a good reason why you didn’t

    • dynomight@lemmy.worldOPMEnglish
      2·
      1 month ago

      Well, Einstein summation is good, but it only does multiplication and sums. (Or, more generally, some scalar operation and some scalar reduction.) I want a notation that works for ANY type of operation, including non-scalar ones, and that’s what DumPy does. So I’d argue it moves further than Einstein summation.

      • ferflo@lemmy.worldEnglish
        1·
        1 month ago

        There’s einx which allows expressing most tensor operations using einsum-like notation: https://github.com/fferflo/einx (Disclaimer: I’m the author). Dumpy and einx actually seem similar to me in that they both use axis names to represent for-loops/vectorization over some simpler, underlying operation.

        • dynomight@lemmy.worldOPMEnglish
          2·
          1 month ago

          Hey, thanks for pointing this out! I quite like the bracket notation for indicating axes that operations should be applied “to” vs. “over”.

          One question I have—is it possible for me as a user to define my own function and then apply it with einx-type notation?

          • ferflo@lemmy.worldEnglish
            1·
            1 month ago

            Thanks! You can use einx.vmap for custom operations:

            def my_dot(x, y):
    return jnp.sum(x * y)

z = einx.vmap("a [c], b [c] -> b a", x, y, op=my_dot)

            Or like so:

            def my_dot(x, y):
    return jnp.sum(x * y)
my_dot = partial(einx.vmap, op=my_dot)

z = my_dot("a [c], b [c] -> b a", x, y)
            • dynomight@lemmy.worldOPMEnglish
              1·
              1 month ago

              OK, I gave it a shot on the initial example in my post:

              import einx
from jax import numpy as jnp
import numpy as onp
import jax

X = jnp.array(onp.random.randn(20,5))
Y = jnp.array(onp.random.randn(30,5))
A = jnp.array(onp.random.randn(20,30,5,5))

def my_op(x,y,a):
    print(x.shape)
    return y @ jnp.linalg.solve(a,x)

Z = einx.vmap("i [m], j [n], i j [m n]->i j", X, Y, A, op=my_op)

              Aaaaand, it seemed to work the first time! Well done!

              I am a little confused though, because if I use "i [a], j [b], i j [c d]->i j" it still seems to work, so maybe I don’t actually 100% understand that bracket notation after all…

              Two more thoughts:

              1. I added a link.
              2. You gotta add def wrap(fun): partial(vmap, op=fun) for easy wrapping. :)
              • ferflo@lemmy.worldEnglish
                1·
                1 month ago

                Thanks for the mention!

                Regarding the naming of axes: einx.vmap doesn’t know anything about my_op, other than that it has the signature "m, n, m n -> " in the first case and "a, b, c d -> " in the second case. Both are valid if you pass the right inputs shapes. You get different behavior for incorrect input shapes though: In the first case, einx will raise an exception before calling my_op due to failing the shape resolution (e.g. due to multiple different values for m). In the second case, einx will assume the shapes to be correct (and it can’t know they aren’t correct before calling my_op), so the error will be raised somewhere in my_op.

                The decorator for einx.vmap is a good point. I did only realize when typing the above comment that wrapping is a nice way of writing the operation in the first place. :D

                • dynomight@lemmy.worldOPMEnglish
                  1·
                  1 month ago

                  Ah, I see, very nice. I wonder if it might make sense to declare the dimensions that are supposed to match once and for all when you wrap the function?

                  E.g. perhaps you could write:

                  @new_wrap('m, n, m n->')
def my_op(x,y,a):
    return y @ jnp.linalg.solve(a,x)

                  to declare the matching dimensions of the wrapped function and then call it with something like

                  Z = my_op('i [:], j [:], i j [: :]->i j', X, Y, A)

                  It’s a small thing but it seems like the matching declaration should be done “once and for all”?

                  (On the other hand, I guess there might be cases where the way things match depend on the arguments…)

                  Edit: Or perhaps if you declare the matching shapes when you wrap the function you wouldn’t actually need to use brackets at all, and could just call it as:

                  Z = my_op('i :, j :, i j : :->i j', X, Y, A)

                  ?

  • sjudubya@lemmy.worldEnglish
    2·
    1 month ago

    Nice suggestions, and I like your API. Using the context manager to look like loops is nifty.

    • dynomight@lemmy.worldOPMEnglish
      11·
      1 month ago

      Thanks, the one problem with that is that you have to use dumpy.wrap if you ever create a function that uses loops and then you want to call it inside another loop. But I don’t see any way around that.