dynomight internet forum

73 readers

1 users here now

founded 10 months ago

MODERATORS

dynomight@lemmy.world

DumPy: NumPy except it’s OK if you’re dum (dynomight.net)

submitted 1 week ago by dynomight@lemmy.world to c/dynomight@lemmy.world

14 comments fedilink hide all child comments

you are viewing a single comment's thread
view the rest of the comments

[–] dynomight@lemmy.world 1 points 1 week ago* (last edited 1 week ago) (1 children)

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)

[–] ferflo@lemmy.world 1 points 1 week ago

I did at some point consider adding a new symbol (like ":") that would act as a new axis with a unique name, but have been hesitant so far, since it adds to the complexity of the notation. There are a bunch of ideas for improving quality-of-life in einx, but so far I've tried erring on the side of less complexity (and there's probably some cases where I should've adhered to this more); to keep a low barrier of entry and also not end up with a mess of many different rules that classical tensor notation is in (you made good points about that here...). There are indeed cases where the operation depends on names in the brackets, e.g. einx.dot("a [b], [b c], [c] d", ...), so the ":" would be an additional variant rather than a simplification.

What I like better about using actual names is also that they convey semantics

einx.softmax("b h q [k]", x) # Ok, softmax along key dimension
einx.softmax("b h q [:]", x) # Softmax along... last dimension?

and indicate corresponding axes in consecutive operations (although this is not enforced strictly).

Defining the shapes of a vmapped operation in the decorator sounds like a good idea. It would probably require a kind of pattern matching to align the inner and outer expression (e.g. to also allow for n-dimensional inputs or variadic arguments to the custom, vmapped function).