[0.5] leaf values are no longer saved by default compared to 0.4

[Butanium]

Issue linked to this forum post: https://discuss.ndif.us/t/nnsight-0-5-prerelease-feedback-requested/64/39?u=clement_dumas

I was surprised that you need to save leaf values like integers or None. This is another breaking change compared to 0.4

from nnsight import LanguageModel

model = LanguageModel("gpt2")
# a = "default"
with model.trace("hi") as tracer:
    a = None
print(f"a: {a}")
> prints default or fails with NameError: name 'a' is not defined

this happened to me with this code failing

    if cache_inputs:
        input_ids = model.input_ids.save()
    else:
        input_ids = None

I will just initialize input_ids to None outside the trace, but this could be a surprising to new user. I could also just do

    if cache_inputs:
        input_ids = model.input_ids
    else:
        input_ids = None
    input_ids.save()

Maybe saving the non-tracing values could be a better default for 0.5? happy to hear your thoughts on this. but e.g.

my_var = 1
my_var = my_var.save()

looks quite verbose and can only be shortened with my_var = int(1).save() but there is no equivalent for None AFAIK

For None an alternative is

def none():
   return None

with model ...:
   a = none().save()

[JadenFiotto-Kaufman]

I could make None default saved. Not sure about other primitives.

[Butanium]

Yeah more generally this behavior of "changes of values on your predefined locals are only taken into account if you save the value" seems error prone. I'm wondering if you could throw a warning if a saved value is overwritten with a non saved value?

Like this also looks tricky:

from nnsight import LanguageModel
import torch as th
model = LanguageModel("gpt2")
# x = "default"
with model.trace("a") as trace:
    x = th.tensor([1, 2, 3]).save()
    x = 2
print(f"{x=}")
> x=default / NameError: name 'x' is not defined