`(zeros|ones|full)_like` and `axes`/`tags`

[majosm]

Argh, accidentally pressed Enter. Ignore for the moment. 🙂 (See below.)

[majosm]

While debugging axis tagging I've run into a couple of snags with zeros_like, etc. and wanted to mention them for discussion.

1. 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?
2. 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?

cc @inducer, @kaushikcfd, @matthiasdiener

[inducer]

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.

[inducer]

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]

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?

[inducer]

I was implicitly using the terminology

after = actx.np.zeros_like(before)

But what I proposed won't work because it would require mutating the tags of before in-place. So forget I said anything.

[inducer]

Two options for navigating this:

• 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.

[inducer]

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.)

[kaushikcfd]

See inducer/pytato#604 for a potential implementation.