Add a Notebook about Trace, BoundSymbol, Symbol and Proxy

[kshitij12345]

No description provided.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:00Z
----------------------------------------------------------------

typos:

into its components.

For Inspection, it...

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:01Z
----------------------------------------------------------------

Line #2.    exec_trace: list[Trace] = thunder.last_traces(jfn)[-1]

It is natural to wonder what the preceding elements in this list are. Maybe a sentence addressing that could be added.

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:01Z
----------------------------------------------------------------

While names are discussed in the later section about Proxies, I expected some treatment of where the t*s come from. Also, the t_*s are completely mysterious as they don't appear in the printed trace.

crcrpar commented on 2024-11-12T12:47:38Z
----------------------------------------------------------------

really good point.

t_* names are temporary and not visible almost always, as they are renamed accordingly to the function signature around IIUC

lightning-thunder/thunder/core/jit_ext.py

Lines 1794 to 1806 in f7b2e15

	# Update prologue trace by renaming proxies which are passed from prologue to the computation trace
	prologue_trace = _apply_trace_proxy_rename(prologue_trace, restrict_proxy_swapmap(pro_to_comp_proxies))
	update_tags(ctx._proxy_swapmap)
	# Update computation trace by renaming proxies which are in the ctx._proxy_swapmap
	computation_trace = _apply_trace_proxy_rename(computation_trace, ctx._proxy_swapmap, "computation")
	# Update epilogue trace by renaming proxies which are passed to the epilogue trace from prologue and computation traces
	if epilogue_trace:
	epilogue_trace = _apply_trace_proxy_rename(
	epilogue_trace, restrict_proxy_swapmap(pro_to_epi_proxies + comp_to_epi_proxies), "epilogue"
	)

kshitij12345 commented on 2024-11-12T12:51:03Z
----------------------------------------------------------------

Interesting, at first glance, I assumed all names were present in the trace. Will try to understand what is happening here. Thanks!

kshitij12345 commented on 2024-11-12T12:57:02Z
----------------------------------------------------------------

Thanks @crcrpar, I typed the above comment before seeing your comment.

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:02Z
----------------------------------------------------------------

I find this section very informative. I like that all the BoundSymbols are put on display. It makes concrete a very abstract class.

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:03Z
----------------------------------------------------------------

typo: its

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:03Z
----------------------------------------------------------------

typos: an operation to be performed on its inputs. It is comprise of an id and a meta function. The meta function

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:04Z
----------------------------------------------------------------

Line #6.    # These inputs and outputs are usually Proxy object (see below)

typo: Proxy objects

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:05Z
----------------------------------------------------------------

Line #8.        # In this tutorial, we will re-use one of the existing symbols that already exists.

This comment is confusing to me. So, thunder.torch is a module that contains Symbols (ie, not functions for computation) and rather than creating a Symbol from scratch in this tutorial, we instead import it from this module? Maybe underscoring the difference between torch and ltorch might drive this point home.

Since add is again used in the next code block, I would find it clearer if this didn't look like a local import.

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:06Z
----------------------------------------------------------------

Line #18.    si = SigInfo("my_trace")

SigInfo comes out of nowhere and I think it would be beneficial to add some explanation here.

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:06Z
----------------------------------------------------------------

Why does this trace look simpler than both the original and the one below?

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

beverlylytle commented on 2024-11-11T13:13:07Z
----------------------------------------------------------------

Line #14.        o = add(x, y)

Why is o being reassigned here? Is it necessary to first create the TensorProxy ? The comment seems to be indicating that by calling add(x, y) within the trace context there are side effects on the context: A BoundSymbol for add is created and appended. Does it not also create and append a BoundSymbol for the output of add ? Is the assignment operator overloaded in some way that also has side effects on the trace? Some of this information is included in the leaf bubble from Symbol in the diagram. I think it would be worthwhile to repeat that text (or something very close to it) again here.

kshitij12345 commented on 2024-11-12T12:48:03Z
----------------------------------------------------------------

Actually, the first o is not needed (I just forgot to remove it). It is correct that add(x, y) will create a BoundSymbol for add to the trace and also correctly set the output. Have updated the code, thank you!

[crcrpar]

t_* names are temporary and not visible almost always, as they are renamed accordingly to the function signature around IIUC

lightning-thunder/thunder/core/jit_ext.py

Lines 1794 to 1806 in f7b2e15

	# Update prologue trace by renaming proxies which are passed from prologue to the computation trace
	prologue_trace = _apply_trace_proxy_rename(prologue_trace, restrict_proxy_swapmap(pro_to_comp_proxies))
	update_tags(ctx._proxy_swapmap)
	# Update computation trace by renaming proxies which are in the ctx._proxy_swapmap
	computation_trace = _apply_trace_proxy_rename(computation_trace, ctx._proxy_swapmap, "computation")
	# Update epilogue trace by renaming proxies which are passed to the epilogue trace from prologue and computation traces
	if epilogue_trace:
	epilogue_trace = _apply_trace_proxy_rename(
	epilogue_trace, restrict_proxy_swapmap(pro_to_epi_proxies + comp_to_epi_proxies), "epilogue"
	)

---

View entire conversation on ReviewNB

[kshitij12345]

Actually, the first o is not needed (I just forgot to remove it). It is correct that add(x, y) will create a BoundSymbol for add to the trace and also correctly set the output. Have updated the code, thank you!

---

View entire conversation on ReviewNB

[kshitij12345]

Interesting, at first glance, I assumed all names were present in the trace. Will try to understand what is happening here. Thanks!

---

View entire conversation on ReviewNB

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

crcrpar commented on 2024-11-12T12:53:06Z
----------------------------------------------------------------

Line #3.    from thunder.core.trace import TraceCtx as Trace, tracectx

out of curiosity, why do we import TraceCtx as Trace in some places?

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

crcrpar commented on 2024-11-12T12:53:07Z
----------------------------------------------------------------

Line #13.        add_bsym = BoundSymbol(sym=add, args=(x, y), kwargs={}, output=o)

I thought we could add_bsym = add.bind(x, y, output=o) as well. I'm curious to hear others' take/preference

kshitij12345 commented on 2024-11-12T12:59:18Z
----------------------------------------------------------------

That should work as well. The reason I choose to construct BoundSymbol was to show-case how it can be directly constructed as we talk about it in the earlier section. We have multiple ways to add a BoundSymbol to a trace (which can be confusing sometimes).

[kshitij12345]

Thanks @crcrpar, I typed the above comment before seeing your comment.

---

View entire conversation on ReviewNB

[kshitij12345]

That should work as well. The reason I choose to construct BoundSymbol was to show-case how it can be directly constructed as we talk about BoundSymbol in the earlier section. We have multiple ways to add a BoundSymbol to a trace (which can be confusing sometimes).

---

View entire conversation on ReviewNB

[t-vi]

@kshitij12345 Great stuff. Thank you for doing this!

re creating BoundSymbols, the recommendation should be:

• use bound_symbol.from_bsym(...) to create a "derivative" (e.g. replacement, but this is not automatic) using most of the things from an existing bound symbol.
• use symbol(*args, **kwargs) if to get a new output (typically a newly created proxy or a tuple of them), - added to the trace automatically,
• use symbol.bind(..., output=...) to use an existing output,
• don't use the constructor BoundSymbol(...).

If you want to have new bound symbols somewhere else then at the end, use tracectx.push_scope([]) before and new_bsyms = tracectx.pop_scope() afterwards.

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

crcrpar commented on 2024-11-13T15:06:56Z
----------------------------------------------------------------

is the left top code snippet the first one of thunder.last_traces , out of curiosity?

I don't demand an update of it but prims.convert_element_type(1, flot) doesn't seem to have effects, so I just am curious.

That said, maybe use the example used below, f(x, y): return x+y, might be nice to do