Pulling an intermediate variable outside a function

[astanziola]

One feature that would be quite nice to have is the ability to retrieve an intermediate variable from a jitted function, along with its results, without needing to modify the function signature. This could be especially useful for nested functions.

My thoughts are along these lines: take, for instance, the functions

@jit
def foo(x):
  y = 2*x
  return x + y

@jit
def bar(x): 
  return foo(x)*x

I would like to have some method, perhaps called pull_variable, that could extract a named_variable like so:

@jit
def foo(x):
  y = named_variable(2*x, "qux")
  return x + y

@pull_variable("qux")
@jit
def bar(x): 
  return foo(x)*x

Now, I would be able to execute:

result, variables = bar(x)

Where variables = {"qux": y}.

While I don't see a reason why this shouldn't be feasible, I'm curious if there's a way to achieve this in Jax currently?

If not, this could indeed be a great feature, and I can foresee quite a few potential use cases for it, such as:

• Inspecting a variable that might be causing an issue in a complex Jax function.
• Monitoring the values of specific variables during large computations for optimization or performance analysis, without the need to store all intermediate computations.
• Performing unit testing on certain internal computations or values within my JIT compiled functions, without having to create separate testing functions.
• Implementing machine learning models where I need to monitor the internal state (like gradients or layer outputs) for debugging or visualization purposes.
• Teaching or demonstrating an algorithm where accessing and explaining intermediate values would improve clarity and learning.
• Running simulations or numerical experiments and wanting to record or analyze the trajectories of specific variables without significantly affecting performance.

On a tangential note, there may be some complimentary arguments for having a push_variable decorator that modifies the input signature of a function, and allows injecting a value in place of a variable in a computational graph, but that's partially beyond the scope of this question 😄

[davisyoshida]

I believe Oryx implements this idea (see here), but when I tried to install it with pip it looks like it may not be compatible with the most recent JAX version.

[astanziola]

That looks exactly like what I was talking about, thanks!