Is "a = a + b" or "a += b" applied in-place inside jit() ?

[Dong-Jiahuan]

According to https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at

inside a jit() compiled function, expressions like x = x.at[idx].set(y) are guaranteed to be applied in-place.

I wanna know that is "a = a + b" or "a += b" applied in-place inside jit() ?

[jakevdp]

This question is hard to answer, because it is based on the assumption that a sub-computation like a += b within a jit-compiled function results in a buffer that stores the computation output. In jit, this is not the case. So the technical anwer to your question is no, this computation is not applied in-place (the input buffer is not modified). But this does not mean that a new buffer is allocated to store the result of a + b. Rather, the compiler will typically fuse computations like this with neighboring computations to avoid the need to store the values at all.

You can get some visibility into the choices the compiler is making for a particular program using ahead-of-time lowering and compilation APIs. For example, here's how the compiler handles a simple version of a += b:

import jax

def f(x):
  x -= 1
  return x.sum()

x = jnp.arange(10.0)
print(jax.jit(f).lower(x).compile().as_text())

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

%region_0.6 (Arg_0.7: f32[], Arg_1.8: f32[]) -> f32[] {
  %Arg_0.7 = f32[] parameter(0)
  %Arg_1.8 = f32[] parameter(1)
  ROOT %add.9 = f32[] add(f32[] %Arg_0.7, f32[] %Arg_1.8), metadata={op_name="jit(f)/jit(main)/reduce_sum[axes=(0,)]" source_file="<ipython-input-9-fbf40f8ff548>" source_line=5}
}

%fused_computation (param_0.2: f32[10]) -> f32[] {
  %param_0.2 = f32[10]{0} parameter(0)
  %constant.3 = f32[] constant(-1)
  %broadcast.1 = f32[10]{0} broadcast(f32[] %constant.3), dimensions={}
  %add.1 = f32[10]{0} add(f32[10]{0} %param_0.2, f32[10]{0} %broadcast.1), metadata={op_name="jit(f)/jit(main)/sub" source_file="<ipython-input-9-fbf40f8ff548>" source_line=4}
  %constant.1 = f32[] constant(0)
  ROOT %reduce.0 = f32[] reduce(f32[10]{0} %add.1, f32[] %constant.1), dimensions={0}, to_apply=%region_0.6, metadata={op_name="jit(f)/jit(main)/reduce_sum[axes=(0,)]" source_file="<ipython-input-9-fbf40f8ff548>" source_line=5}
}

ENTRY %main.11 (Arg_0.1: f32[10]) -> f32[] {
  %Arg_0.1 = f32[10]{0} parameter(0), sharding={replicated}
  ROOT %fusion = f32[] fusion(f32[10]{0} %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(f)/jit(main)/reduce_sum[axes=(0,)]" source_file="<ipython-input-9-fbf40f8ff548>" source_line=5}
}

Notice that in the compiled output, the entire operation (x - 1).sum() appears within a fusion, which essentially means the compiler will compute it in one pass without explicitly storing intermediate values.

So the answer to your question is no, this operation is not applied in-place (as in, the input buffer is not modified) but neither are the values copied to a different buffer. Does that help answer your question?