You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
While debugging axis tagging I've run into a couple of snags with zeros_like, etc. and wanted to mention them for discussion.
In Copy axes and tags in zeros_like/full_like聽#229, we modified these functions to propagate both axes and tags. In retrospect, I'm not sure if tags should be propagated. For example, when using Add tag to store array creation traceback聽pytato#284, this caused the CreatedAt tag from the source array to be propagated to the new array of zeros, giving the impression that the array was created somewhere where it wasn't. I'm not terribly familiar with the other tags, but I'm wondering if it makes much sense to propagate those either?
Another issue I ran into is that since the result of zeros_like does not explicitly depend on the incoming source array, AFAIK axes are not able to be inferred from this array. In the case of a simulation state initialization, where the computation being done might only depend on nodes (i.e., no grudge operators are used), removing the dependency on the source could cut off the only source of axis inferral for the new array. I don't have a clear idea of how to fix this, apart from using the old 0*array method. Thoughts?
These are good questions. Thanks for raising them!
For your first question, a simple-but-perhaps-too-ad-hoc way of addressing the issue would be to tell tags (e.g. by being a subclass of PropagatableTag or so, with TBD semantics) whether they should propagate in inference and in situations like zeros_like. "Descriptive" tags definitely should propagate, but many others shouldn't. I'd welcome opinions on this.
The second one is thornier. 馃 The problem, IIUC, is that information propagates from the "source" array to the "target" array only at the moment when zeros_like runs, but not after that. There's a menu of (bad?) options here for how to ensure information continues to propagate:
We could make a tag 馃う to say so. A problem with this is that such a tag would hold a reference to another array, which would have to be replaced upon every DAG rewrite. I would guess that this would drive up the cost of DAG rewrites, because suddenly we have to go hunt for these arrays in the tags on every array we traverse.
We could make a less-bad version of 0*array. One way to do so would be a function zero(...) that returns zero no matter the input. (This would address the issue that 0*array propagates NaNs.) This would allow the connection from source to target to exist just long enough for tag propagation to work (with a definable axis mapping). Upon code generation, we might need to find a way to drop the spurious dependency edge.
None of these options are particularly lovely IMO. If there are better ideas, I'd be all ears.
I think I just thought of a possible solution for the second problem. Rather than stick a reference to another array into a tag, we can simply assign a unique (e.g. by incrementing a number) identifier to the axes in the "before" and "after" arrays (... also via a tag). Tag propagation could then create equations for matching identifiers. (see below)
I think I just thought of a possible solution for the second problem. Rather than stick a reference to another array into a tag, we can simply assign a unique (e.g. by incrementing a number) identifier to the axes in the "before" and "after" arrays (... also via a tag). Tag propagation could then create equations for matching identifiers.
I forgot to ask about this this morning. I'm not quite following the "before"/"after", can you explain a bit more?
First, maybe "tag propagation" is overthinking it. I can't think of a reason why we can't assign axis tags eagerly. Then we could just copy the tags during zeros_like and be done with it. It sure is easier to debug. Axis labeling is a core operating principle behind https://docs.xarray.dev/en/stable/index.html, and people seem to like it OK. Array creation mostly happens under infrastructure control (discr.zeros(), grudge operators). A downside is that we'd have to hardcode the propagation logic in many places, so it would become basically impossible to change.
If it turns out that tag propagation is needed after all, we can just unconditionally assign "every axis under the sun" a name and just refer to it. The one complication I can think of with respect to this is pickling/unpickling, but this is manageable.
Another option, from a discussion with @kaushikcfd: Define a (scalar) function zero(x) and have zeros_like emit an IndexLambda with zeros(a[_0,_1]) as an expression. This will appropriately propagate tags with the logic as it exists now, and a pass to replace all the zeros(...) with literal zeros (and eliminating extraneous dependencies) is straightforward to write.
Thus far, I'm gravitating towards this solution. (As a reminder, 0*a[_0,_1] is not usable because it needs to propagate NaNs for correctness, so a needs to still be computed, and it remains an input.)
Activity
[-]`(zeros_like` and [/-][+]`(zeros|ones|full)_like` and `axes`/`tags`[/+]majosm commentedon May 23, 2023
While debugging axis tagging I've run into a couple of snags with
zeros_like, etc. and wanted to mention them for discussion.axesandtagsinzeros_like/full_like聽#229, we modified these functions to propagate bothaxesandtags. In retrospect, I'm not sure iftagsshould be propagated. For example, when using Add tag to store array creation traceback聽pytato#284, this caused theCreatedAttag from the source array to be propagated to the new array of zeros, giving the impression that the array was created somewhere where it wasn't. I'm not terribly familiar with the other tags, but I'm wondering if it makes much sense to propagate those either?zeros_likedoes not explicitly depend on the incoming source array, AFAIK axes are not able to be inferred from this array. In the case of a simulation state initialization, where the computation being done might only depend onnodes(i.e., no grudge operators are used), removing the dependency on the source could cut off the only source of axis inferral for the new array. I don't have a clear idea of how to fix this, apart from using the old0*arraymethod. Thoughts?cc @inducer, @kaushikcfd, @matthiasdiener
inducer commentedon May 23, 2023
These are good questions. Thanks for raising them!
For your first question, a simple-but-perhaps-too-ad-hoc way of addressing the issue would be to tell tags (e.g. by being a subclass of
PropagatableTagor so, with TBD semantics) whether they should propagate in inference and in situations likezeros_like. "Descriptive" tags definitely should propagate, but many others shouldn't. I'd welcome opinions on this.The second one is thornier. 馃 The problem, IIUC, is that information propagates from the "source" array to the "target" array only at the moment when
zeros_likeruns, but not after that. There's a menu of (bad?) options here for how to ensure information continues to propagate:0*array. One way to do so would be a functionzero(...)that returns zero no matter the input. (This would address the issue that0*arraypropagates NaNs.) This would allow the connection from source to target to exist just long enough for tag propagation to work (with a definable axis mapping). Upon code generation, we might need to find a way to drop the spurious dependency edge.None of these options are particularly lovely IMO. If there are better ideas, I'd be all ears.
temporarily use old bad method of creating zeros/ones/full
inducer commentedon May 1, 2025
I think I just thought of a possible solution for the second problem. Rather than stick a reference to another array into a tag, we can simply assign a unique (e.g. by incrementing a number) identifier to the axes in the "before" and "after" arrays (... also via a tag). Tag propagation could then create equations for matching identifiers.(see below)majosm commentedon May 2, 2025
I forgot to ask about this this morning. I'm not quite following the "before"/"after", can you explain a bit more?
inducer commentedon May 5, 2025
I was implicitly using the terminology
But what I proposed won't work because it would require mutating the tags of
beforein-place. So forget I said anything.inducer commentedon May 16, 2025
Two options for navigating this:
zeros_likeand be done with it. It sure is easier to debug. Axis labeling is a core operating principle behind https://docs.xarray.dev/en/stable/index.html, and people seem to like it OK. Array creation mostly happens under infrastructure control (discr.zeros(), grudge operators). A downside is that we'd have to hardcode the propagation logic in many places, so it would become basically impossible to change.inducer commentedon May 16, 2025
Another option, from a discussion with @kaushikcfd: Define a (scalar) function
zero(x)and havezeros_likeemit anIndexLambdawithzeros(a[_0,_1])as an expression. This will appropriately propagate tags with the logic as it exists now, and a pass to replace all thezeros(...)with literal zeros (and eliminating extraneous dependencies) is straightforward to write.Thus far, I'm gravitating towards this solution. (As a reminder,
0*a[_0,_1]is not usable because it needs to propagate NaNs for correctness, soaneeds to still be computed, and it remains an input.)kaushikcfd commentedon May 23, 2025
See inducer/pytato#604 for a potential implementation.