xla opcode

[never-to-never]

xla_comp = jax.xla_computation(fun)(inputs)
xla_text = xla_comp.as_serialized_hlo_module_proto()

from tensorflow.compiler.xla.service import hlo_pb2
module = hlo_pb2.HloModuleProto() module.ParseFromString(xla_text)

The module obtained in this way contains many instructions, and each instruction corresponds to an opcode. Where can I get all the opcode lists used by the jax program?

[jakevdp]

Can you clarify your question? Are you asking how you can access the list of all opcodes generated by a particular JAX program, or are you asking for more general documentation on all the opcodes that can possibly be generated by any JAX program?

[never-to-never]

@jakevdp hello, What I want to ask is a list of all opcodes that may be generated by any jax program?
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/hlo/ir/hlo_opcode.h
I saw a HLO_OPCODE_LIST from the source code of tensorflow, I don't know if this is the list of all opcodes that the jax program may generate?

[jakevdp]

The full list of opcodes can be found in the OpenXLA repository, here: https://github.com/openxla/xla/blob/594fdea6c79a0e93386556eaa5958608698e5ed2/xla/hlo/ir/hlo_opcode.h#L48