Which dataclass members become HLO parameters? #30251

staticimport · 2025-07-16T14:42:24Z

staticimport
Jul 16, 2025

Below is a toy example that has a dataclass containing 2 members, but my function only uses one of them. As part of compiling down to HLO, the unused member is eliminated. How can I programmatically tell which parameters the HLO module wants?

Motivation: I'm working on using XLA AOT for C++ inference, so in C++ I need to know exactly what inputs to pass in.

import jax, jax.numpy as jnp
from flax import struct
@struct.dataclass

class Foo:
    bar: jnp.ndarray
    baz: jnp.ndarray

def fn(foo: Foo):
    return jnp.sum(foo.baz)

foo = Foo(jnp.arange(10), jnp.arange(10))
hlo = jax.jit(fn).lower(foo).compiler_ir("hlo")
print(hlo.as_hlo_text())

Output below is completely identical whether fn( ) above sums foo.bar or foo baz:

HloModule jit_fn, entry_computation_layout={(s32[10]{0})->s32[]}, frontend_attributes={xla.sdy.meshes={}}

region_0.3 {
  Arg_0.4 = s32[] parameter(0)
  Arg_1.5 = s32[] parameter(1)
  ROOT add.6 = s32[] add(Arg_0.4, Arg_1.5)
}

ENTRY main.8 {
  Arg_0.1 = s32[10]{0} parameter(0)
  constant.2 = s32[] constant(0)
  ROOT reduce.7 = s32[] reduce(Arg_0.1, constant.2), dimensions={0}, to_apply=region_0.3
}

Answered by jakevdp

Jul 16, 2025

I don't know how to get this information in the compiler_ir representation, but if you print .compile().as_text() then you get an HLO representation that includes the Python provenance of each variable:

lowered = jax.jit(fn).lower(foo).compile()
print(lowered.as_text())

HloModule jit_fn, is_scheduled=true, entry_computation_layout={(s32[10]{0})->s32[]}, allow_spmd_sharding_propagation_to_parameters={true}, allow_spmd_sharding_propagation_to_output={true}

%region_0.3 (Arg_0.4: s32[], Arg_1.5: s32[]) -> s32[] {
  %Arg_0.4 = s32[] parameter(0), metadata={op_name="jit(fn)/jit(main)/reduce_sum"}
  %Arg_1.5 = s32[] parameter(1), metadata={op_name="jit(fn)/jit(main)/reduce_sum"}
  ROOT %add.6 =…

View full answer

jakevdp · 2025-07-16T15:41:25Z

jakevdp
Jul 16, 2025
Maintainer

I don't know how to get this information in the compiler_ir representation, but if you print .compile().as_text() then you get an HLO representation that includes the Python provenance of each variable:

lowered = jax.jit(fn).lower(foo).compile()
print(lowered.as_text())

HloModule jit_fn, is_scheduled=true, entry_computation_layout={(s32[10]{0})->s32[]}, allow_spmd_sharding_propagation_to_parameters={true}, allow_spmd_sharding_propagation_to_output={true}

%region_0.3 (Arg_0.4: s32[], Arg_1.5: s32[]) -> s32[] {
  %Arg_0.4 = s32[] parameter(0), metadata={op_name="jit(fn)/jit(main)/reduce_sum"}
  %Arg_1.5 = s32[] parameter(1), metadata={op_name="jit(fn)/jit(main)/reduce_sum"}
  ROOT %add.6 = s32[] add(s32[] %Arg_0.4, s32[] %Arg_1.5), metadata={op_name="jit(fn)/jit(main)/reduce_sum" source_file="/tmp/ipython-input-7-3427796640.py" source_line=10}
}

ENTRY %main.8 (Arg_0.1: s32[10]) -> s32[] {
  %Arg_0.1 = s32[10]{0} parameter(0), metadata={op_name="foo.baz"}
  %constant.2 = s32[] constant(0)
  ROOT %reduce.7 = s32[] reduce(s32[10]{0} %Arg_0.1, s32[] %constant.2), dimensions={0}, to_apply=%region_0.3, metadata={op_name="jit(fn)/jit(main)/reduce_sum" source_file="/tmp/ipython-input-7-3427796640.py" source_line=10}
}

(Notice the metadata={op_name="foo.baz"} in the main block).

3 replies

staticimport Jul 16, 2025
Author

Ah lovely, thanks so much! Only now read in the doc string that compiler_ir() is just for debugging? But interestingly it lacks this metadata, and it still has compiled/executed successfully in PjRt for me.

This still feels a bit janky, having to regex parse the HLA, vs. some python object telling me which parameters were used. I noticed when running this for my real code, the metadata doesn't necessarily show up on the parameter(X) line, e.g.

%param_2.21 = s8[1,1,1]{2,1,0} parameter(2)
%bitcast.44 = s8[1,1]{1,0} bitcast(s8[1,1,1]{2,1,0} %param_2.21), metadata={op_name="dataset.data.alias_idx"}

Still, at least it's feasible to parse this. Thanks again!

jakevdp Jul 16, 2025
Maintainer

In general this is not required – the normal JAX call/lowering pipeline takes care of all parameter mapping for you. The as_hlo_text() and related methods are mainly provided for debugging purposes, not necessarily meant to be used in execution pipelines, so it's not surprising that using the code in an unintended way requires some hacks.

staticimport Jul 16, 2025
Author

Yeah fair enough. Not my goal to be doing things in unintended ways, not this time anyway. 😅 I'm just trying to export my model for C++ (cpu or gpu) inference, went the route of XLA AOT with the PjRT C Api instead of jax2tf or jax2onnx. Maybe those are cleaner? I didn't find documentation on inference best practices, so admittedly stumbling in the dark a bit.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Which dataclass members become HLO parameters? #30251

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Replies: 1 comment 3 replies

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Select a reply

Uh oh!

Which dataclass members become HLO parameters? #30251

Uh oh!

Uh oh!

staticimport Jul 16, 2025

Replies: 1 comment · 3 replies

Uh oh!

Uh oh!

jakevdp Jul 16, 2025 Maintainer

Uh oh!

staticimport Jul 16, 2025 Author

Uh oh!

jakevdp Jul 16, 2025 Maintainer

Uh oh!

Uh oh!

staticimport Jul 16, 2025 Author

staticimport
Jul 16, 2025

Replies: 1 comment 3 replies

jakevdp
Jul 16, 2025
Maintainer

staticimport Jul 16, 2025
Author

jakevdp Jul 16, 2025
Maintainer

staticimport Jul 16, 2025
Author