How to understand tracing and static variables in Jax

[qsh-zh]

I am learning jax. However, I find it is difficult to understand some concepts related with tracing in the tutorial.

1. What is tracing when we discuss tracing and static variables. I only encountered the trace when using logger. Obviously, it has a different meaning in jax programming language. My naive understanding is by saying tracing something, we do some compiling and optimization to speed up by providing more information.

2. Why do we need tracing? And how can we distinguish the tracing and static variables? I notice that example code here

from functools import partial

@partial(jit, static_argnums=(1,))
def f(x, neg):
  return -x if neg else x

f(1, True)

Since the control flow is determined by the content of the neg, we need to use static variables. My questions is, can I say only when the variables influence the control flow I should use static variables? And does using a lot of static variables slows down the execution?

[qsh-zh]

I found the explanation in Tensorflow are very good for understanding trace.

In TF

• The real computation is done in tf.Graph
• In some sense, tf.Graph is static, requires inputs to have a specified data type and shape
• tracing with tf.function helps bridge python function and tf.Graph
• Execution has two stages, first one is known as tracing with creating a new tf.Graph, the operation is deferred, the second stage runs the calculation

[YouJiacheng]

Tracing is recording the computation with zero FLOP.
Pytorch has tracing too. And the tensor centric dynamic graph of Pytorch use a "Tape" technique, which can be considered as Tracing with FLOPs.

[YouJiacheng]

Actually lax.cond can be used to create a dynamic control flow, by tracing the both branch, which can potentially cause exponential tracing time!
BTW, the inputs determining the shape of array in computation must be static.