Ways to name constants when dumping the hlo module via xla_computation?

[louis-shao]

Hi, when using jax.xla_computation(my_function)(my_args).as_serialized_hlo_module_proto() to dump a function to hlo module (for execution in C++), is it possible to name some of the constants in the proto? I see that some of the instructions in the proto have a metadata field that records some information about how the instruction is generated, so I wonder if it is possible for naming constants too, something like tf.identity(x, name="my_readable_name").

My use case is that I have some data that need to be prepared dynamically in the runtime, but I don't want them to be fed as inputs to the function in every function call, which could add more inference time (which is important to us). I'm hoping that if the constants are named, I can have a placeholder constant in the function and substitute the values in the proto manually before loading the function for execution.

I tried jax.named_scope or jax.named_call around jnp.array(my_numpy_array) but it doesn't keep the names anyhow.

[froystig]

We don't have a mechanism for annotating constants directly. Would it work for you to name a constant function? I'm thinking along the following lines:

import jax

def g(x):
  @jax.jit
  def my_const(): return 3.

  return x + my_const()

>>> print(jax.jit(g).lower(1.).as_text())
module @jit_g attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<f32> {mhlo.sharding = "{replicated}"}) -> (tensor<f32> {jax.result_info = ""}) {
    %0 = call @my_const() : () -> tensor<f32>
    %1 = stablehlo.add %arg0, %0 : tensor<f32>
    return %1 : tensor<f32>
  }
  func.func private @my_const() -> tensor<f32> {
    %0 = stablehlo.constant dense<3.000000e+00> : tensor<f32>
    return %0 : tensor<f32>
  }
}

[louis-shao]

Thank you Roy! This does make it identifiable. Will the function be optimized away after I pass the hlo module to the execution client?

[froystig]

The compiler may choose to optimize it away. I believe that the same goes for uses of named_call.