search

inducer / arraycontext

Choose your favorite numpy-workalike!
6 stars 11 forks source link

Numpy array context rebased #190

Closed inducer closed 1 month ago

inducer commented 2 years ago

This does a fair bit of stuff. The main components are a numpy array context and an in-depth rework of dataclass arraycontainer broadcasting. There is also a tweak to array container serialization ("Container serialization: iterable -> sequence, plus type aliases") that makes the serialized container a sequence and introduces a type alias for it. I've flagged some things that I'm less sure about in a self-review below.

The whole thing is probably best read commit-by-commit to avoid mixing changes that belong to different topic areas.

Closes #93 Closes #235 Closes #95

TODO

inducer commented 2 months ago

Hi all, please take a look, particularly with an eye towards how much you like the direction this is going. It should not break compatibility right now, but it contains some things that will need changing once we remove what's being deprecated now.

inducer commented 2 months ago

Also, sorry about the breakage from my last-minute changes, I'll fix that right away.

alexfikl commented 2 months ago

Was just googling around and found a jax pytree version of container math: https://github.com/google/tree-math It seems to just allow broadcasting scalars at the leaf level, so it's a lot smaller in scope.

alexfikl commented 2 months ago

Following a meeting with @inducer on Friday, he suggested I write down any random thoughts here, so that we can discuss. So here goes before I forget again!

My question was mostly related to what broadcasting we actually need to do? The broadcasting implementation here seems to be very general and very flexible (and quite explicit as well).

Having experience with the pytential part of the stack (as opposed to grudge), these operations seem to be required:

  1. Just direct adding of two DOFArrays (no broadcasting here), e.g. x + y.
    • This should just work no questions asked.
  2. Broadcast DOFArray with a scalar (of some sort? device or host), e.g. x + 1.
    • Can work with a rec_[multi]map_array_container, but having the overload work is nicer.
  3. Broadcast a numpy object array (of scalars or of DOFArray) with a DOFArray, e.g. velocity - mean.
    • This should also work with a [multi]map_array_container, but having the overload work is nicer. Object arrays already broadcast constants, so that will just work.

i.e. all those could be handled (maybe a bit verbosely) without any broadcasting support in arraycontext.with_container_arithmetic. Broadcasting a scalar seems like a common enough case to implement here though.

What fancier broadcasting would we need to support? I guess mirgecom has the most advanced needs here.

alexfikl commented 2 months ago

For what it's worth, tried running more downstreams (mostly for curiosity)

inducer commented 2 months ago

From the tests, the Bcast syntax seems a bit clunky with all the wrapping. Maybe we can introduce some helper functions?

result = broadcast(operator.op, container1, container2, level_or_actx=1)

? Will ponder a bit 😁

Not sure.

I agree the internals are a bit clunky, but I personally find the user-facing bit OK to look at. If we were going the helper function route, we would need at least a "side" in the function name (bcast_left_n_levels?). I'm also not sure that "levels" and "until actx array" are the only rules, so there'd be one helper per rule (not a huge fan of rolling them into one, as level_or_actx suggests.) Having to fish out an operator function is also not super pretty, and allowing "+" would come at a cost.

I may try coding this up to see how it comes out.

inducer commented 2 months ago

What fancier broadcasting would we need to support? I guess mirgecom has the most advanced needs here.

Good question. Grudge can get by without any use of the Bcast stuff. I think I am coming around to not introducing Bcast right now, and to see what we even actually need.

alexfikl commented 2 months ago

Not sure.

I agree the internals are a bit clunky, but I personally find the user-facing bit OK to look at. If we were going the helper function route, we would need at least a "side" in the function name (bcast_left_n_levels?). I'm also not sure that "levels" and "until actx array" are the only rules, so there'd be one helper per rule (not a huge fan of rolling them into one, as level_or_actx suggests.) Having to fish out an operator function is also not super pretty, and allowing "+" would come at a cost.

I may try coding this up to see how it comes out.

I mostly meant it as a purely convenience function for whatever ends up being "a common case". The driving idea was to match something like reduce or all our map_array_container type functions that sort of look like that.

That can go in a future PR, if the need arises, so I don't feel very strongly about it.

inducer commented 2 months ago

Alright, Bcast is gone from this, now in #280 should we need it.