Annotating JAXPR with source code information

[rturrado]

For an input file blah.py like:

import jax

def f(x: float, y: float):
    return x * y

print(jax.make_jaxpr(f)(0.3, 0.4))

The output I get when running python blah.py is:

{ lambda ; a:f32[] b:f32[]. let
    c:f32[] = mul a b
    in (c,) }

Is there any way to keep the original variable names, e.g. x and y, in the JAXPR? Something like a map of invars to source code vars?

Whenever possible, for example for parameter names, as I understand that local variables may be optimised out.

[Qazalbash]

Don't know about jaxpr, but you can annotate HLO with jax.named_scope context manager.

import jax


@jax.jit
def f(x: float, y: float):
    return x * y


@jax.jit
@jax.named_scope("my_named_scope")
def f_with_named_scope(x: float, y: float):
    return x * y


print("Without named scope:")
print(f.lower(0.3, 0.4).compile().as_text())

print("With named scope:")
print(f_with_named_scope.lower(0.3, 0.4).compile().as_text())

Without named scope:
HloModule jit_f, is_scheduled=true, entry_computation_layout={(f32[], f32[])->f32[]}, allow_spmd_sharding_propagation_to_parameters={true,true}, allow_spmd_sharding_propagation_to_output={true}

ENTRY %main.4 (Arg_0.1: f32[], Arg_1.2: f32[]) -> f32[] {
  %Arg_0.1 = f32[] parameter(0), metadata={op_name="x"}
  %Arg_1.2 = f32[] parameter(1), metadata={op_name="y"}
  ROOT %multiply.3 = f32[] multiply(%Arg_0.1, %Arg_1.2), metadata={op_name="jit(f)/jit(main)/mul" source_file="/home/test.py" source_line=6}
}


With named scope:
HloModule jit_f_with_named_scope, is_scheduled=true, entry_computation_layout={(f32[], f32[])->f32[]}, allow_spmd_sharding_propagation_to_parameters={true,true}, allow_spmd_sharding_propagation_to_output={true}

ENTRY %main.4 (Arg_0.1: f32[], Arg_1.2: f32[]) -> f32[] {
  %Arg_0.1 = f32[] parameter(0), metadata={op_name="x"}
  %Arg_1.2 = f32[] parameter(1), metadata={op_name="y"}
  ROOT %multiply.3 = f32[] multiply(%Arg_0.1, %Arg_1.2), metadata={op_name="jit(f_with_named_scope)/jit(main)/my_named_scope/mul" source_file="/home/test.py" source_line=12}
}

See the difference for the context manager name,

HloModule jit_f, is_scheduled=true, entry_computation_layout={(f32[], f32[])->f32[]}, allow_spmd_sharding_propagation_to_parameters={true,true}, allow_spmd_sharding_propagation_to_output={true}

ENTRY %main.4 (Arg_0.1: f32[], Arg_1.2: f32[]) -> f32[] {
  %Arg_0.1 = f32[] parameter(0), metadata={op_name="x"}
  %Arg_1.2 = f32[] parameter(1), metadata={op_name="y"}
-  ROOT %multiply.3 = f32[] multiply(%Arg_0.1, %Arg_1.2), metadata={op_name="jit(f)/jit(main)/mul" source_file="/home/test.py" source_line=6}
+  ROOT %multiply.3 = f32[] multiply(%Arg_0.1, %Arg_1.2), metadata={op_name="jit(f_with_named_scope)/jit(main)/my_named_scope/mul" source_file="/home/test.py" source_line=12}
}

[rturrado]

Thanks @Qazalbash.

I was looking for something more transparent. So you wouldn't need to modify the code, maybe just modify the build flags or some configuration.

Also, considering your example, I am more interested in those op_name=x and op_name=y metadata. What I'm finally aiming for is to have something similar to those metadata in the MLIR, a loc("x") and loc("y") that could be used, together with mlir-use-nameloc-as-prefix, to output generated code such as:

  ROOT %multiply.1 = f32[] multiply(%x, %y)

[jakevdp]

Is there any way to keep the original variable names, e.g. x and y, in the JAXPR?

This is not possible in general; for example:

def f(x):
  for i in range(5):
    x += 1
  return x
print(jax.make_jaxpr(f)(1))

{ lambda ; a:i32[]. let
    b:i32[] = add a 1:i32[]
    c:i32[] = add b 1:i32[]
    d:i32[] = add c 1:i32[]
    e:i32[] = add d 1:i32[]
    f:i32[] = add e 1:i32[]
  in (f,) }

If we were trying to replicate the Python variable names, every jaxpr variable in this expression would have to be named x.

Now, you could imagine various workaround for this (e.g. call them x_1, x_2, etc.) but then you need workarounds for your workarounds (what if there is already a variable called x_1?)

The current implementation has the benefit of simplicity, even if it may be harder to map back to the original code.

[hawkinsp]

Just in case you're not aware: each jaxpr equation already has a full Python traceback attached. We don't display it by default because it's somewhat verbose.

But the variable names are not something we currently attempt to guess (and it wouldn't be possible to do reliably in general).

In [23]: import jax
    ...:
    ...: def f(x: float, y: float):
    ...:     return x * y
    ...:
    ...: print(jax.make_jaxpr(f)(0.3, 0.4).pretty_print(source_info=True))
{ lambda ; a:f32[] b:f32[]. let c:f32[] = mul a b in (c,) } # <ipython-input-23-d4e1599e316a>:4:11 (f)