Does jax.vmap understand not to repeat scalar-only-derived computations?

[jadball]

Given this MWE:

import jax
import jax.numpy as jnp

def get_expensive_scalar(arg1, arg2):
    """Expensive computation that eventually returns a scalar"""
    return arg1 + arg2

def _add_scalar(arr_elem, arg1, arg2):
    """The function we will vmap"""
    scalar = get_expensive_scalar(arg1, arg2)  # does this run for each element of `vec_in` ?
    result = arr_elem + scalar
    return result

add_scalar = jax.jit(jax.vmap(_add_scalar, in_axes=[0, None, None]))

vec_in = jnp.ones(100)

add_scalar(vec_in, 1, 2)

Does jax understand that scalar only needs to be computed once, (because it only depends on scalar inputs) before adding it to the vectorised input?

[jakevdp]

The most definitive way to answer this is to use Ahead of time compilation in order to print the compiled HLO for your function. For example:

print(add_scalar.lower(vec_in, 1, 2).compile().as_text())

HloModule jit__add_scalar, is_scheduled=true, entry_computation_layout={(f32[100]{0}, s32[], s32[])->f32[100]{0}}, allow_spmd_sharding_propagation_to_parameters={true,true,true}, allow_spmd_sharding_propagation_to_output={true}

%fused_computation (param_0: f32[100], param_1.3: s32[], param_2: s32[]) -> f32[100] {
  %param_0 = f32[100]{0} parameter(0)
  %param_1.3 = s32[] parameter(1)
  %param_2 = s32[] parameter(2)
  %add.2 = s32[] add(%param_1.3, %param_2), metadata={op_name="jit(_add_scalar)/vmap()/add" source_file="/Users/vanderplas/github/google/jax/tmp.py" source_line=6 source_end_line=6 source_column=11 source_end_column=22}
  %convert_element_type.0 = f32[] convert(%add.2), metadata={op_name="jit(_add_scalar)/vmap()/convert_element_type" source_file="/Users/vanderplas/github/google/jax/tmp.py" source_line=11 source_end_line=11 source_column=13 source_end_column=30}
  %add.1 = f32[100]{0} broadcast(%convert_element_type.0), dimensions={}, metadata={op_name="jit(_add_scalar)/vmap()/add" source_file="/Users/vanderplas/github/google/jax/tmp.py" source_line=11 source_end_line=11 source_column=13 source_end_column=30}
  ROOT %add.0 = f32[100]{0} add(%param_0, %add.1), metadata={op_name="jit(_add_scalar)/vmap()/add" source_file="/Users/vanderplas/github/google/jax/tmp.py" source_line=11 source_end_line=11 source_column=13 source_end_column=30}
}

ENTRY %main.1 (arr_elem.1: f32[100], arg1.1: s32[], arg2.1: s32[]) -> f32[100] {
  %arr_elem.1 = f32[100]{0} parameter(0), metadata={op_name="arr_elem"}
  %arg1.1 = s32[] parameter(1), metadata={op_name="arg1"}
  %arg2.1 = s32[] parameter(2), metadata={op_name="arg2"}
  ROOT %broadcast_add_fusion = f32[100]{0} fusion(%arr_elem.1, %arg1.1, %arg2.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(_add_scalar)/vmap()/add" source_file="/Users/vanderplas/github/google/jax/tmp.py" source_line=11 source_end_line=11 source_column=13 source_end_column=30}
}

This is admittedly a bit tricky to read, but you can see in this line

%add.2 = s32[] add(%param_1.3, %param_2), metadata={op_name="jit(_add_scalar)/vmap()/add"

that the scalars are added only once before the result is broadcast to the larger shape to be added to the batched input.

There is one detail here: this part of the computation takes place within a fusion, meaning that the compiler may rearrange operations if it is deemed advantageous to do so; that said if your operation really is something expensive (where add is just a stand-in) the compiler would likely treat it differently than it will treat something simple like add.

I hope that helps!

[jadball]

Thanks! In my case, I'm assembling a matrix from some (static) scalar inputs, then performing a dot product over the batched input with the matrix. Hopefully it is still the case there.

[jakevdp]

I suspect it is; you could pass the actual workflow through the AOT compilation to see.