Closed inducer closed 1 month 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.
Also, sorry about the breakage from my last-minute changes, I'll fix that right away.
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.
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:
DOFArray
s (no broadcasting here), e.g. x + y
.
DOFArray
with a scalar (of some sort? device or host), e.g. x + 1
.
rec_[multi]map_array_container
, but having the overload work is nicer.numpy
object array (of scalars or of DOFArray
) with a DOFArray
, e.g. velocity - mean
.
[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.
For what it's worth, tried running more downstreams (mostly for curiosity)
NumpyArrayContext
could work in meshmode (unlikely).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.
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.
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, aslevel_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.
Alright, Bcast
is gone from this, now in #280 should we need it.
This does a fair bit of stuff. The main components are a
numpy
array context and an in-depth rework ofdataclass
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
np.minimum
,np.maximum
missing?