Rehydrating a Jaxpr from HLO / proto?

[jkr26]

Hi JAX community!

Suppose that I can get my hands on an HLO module (proto) representing another author's serialized code, along with any appropriate metadata I might require. Suppose that it is also my desire to 'rehydrate' this logic and use with Jax's function transformations (like autodiff, or even more simply Python function composition). I can find plenty of examples of going from Python to Jaxpr to XLA / HLO to compiled chunks; is there a known way to go the other way around? If not, is there an alternate serialization format that would make such a thing trivially possible?

[mattjj]

Thanks for the question!

You can write an HLO interpreter which walks the HLO program and then evaluates each HLO operation in terms of jax.numpy, or even easier jax.lax (since that's almost 1:1 with HLO). Then you can jax.jit it, apply jax.grad to it. It'd look kind of like onnx2xla.py.

Actually, we wrote something like this years ago. I'll try to find it...

[mattjj]

This is ancient and probably doesn't run anymore. Plus it may not have even been good back then; it was written as a live demo to impress DeepMinders.

import operator as op
from collections import namedtuple
import numpy as onp

from jax.lib.xla_bridge import xla_client
from jax import lax
from jax import core

def _map(f, *args):
  return list(map(f, *args))

etype_names = {
    getattr(xla_client.PrimitiveType, etype_name): etype_name.lower() + 's'
    for etype_name in ["S8", "S16", "S32", "S64", "U8", "U16", "U32", "U64",
                       "F16", "F32", "F64", "C64", "C128", "PRED"]
}

def parameters(computation):
  parameter_instrs = [(instr.parameter_number, instr.id)
                      for instr in computation.instructions
                      if instr.opcode == "parameter"]
  return _map(op.itemgetter(1), sorted(parameter_instrs, key=op.itemgetter(0)))

def constvars(computation):
  return [instr.id for instr in computation.instructions
          if instr.opcode == "constant"]

def constvals(computation):
  def parse_literal(instr):
    etype = int(instr.literal.shape.element_type)
    dtype = xla_client.XLA_ELEMENT_TYPE_TO_DTYPE[etype]
    data = getattr(instr.literal, etype_names[etype])
    shape = instr.literal.shape.dimensions
    return onp.array(data, dtype).reshape(shape)

  return [parse_literal(instr) for instr in computation.instructions
          if instr.opcode == "constant"]

def eqns(computation):
  return [instr for instr in computation.instructions
          if instr.opcode not in {"constant", "parameter"}]

def get_computation(computations, id):
  for c in computations:
    if c.id == id:
      return c
  else:
    assert False

def _slice_indices(instr):
  return zip(*[(si.start, si.limit, si.stride) for si in instr.slice_dimensions])

def debug(*args, **kwargs):
  args, kwargs
  import ipdb; ipdb.set_trace()

standard_ops = {
    "sine": lax.sin,
    "cosine": lax.cos,
    "add": lax.add,
    "subtract": lax.sub,
    "multiply": lax.mul,
    "divide": lax.div,
    "power": lax.pow,
    "select": lax.select,
    "and": lax.bitwise_and,
    "log": lax.log,
    "exponential": lax.exp,
    "negate": lax.neg,
    "maximum": lax.max,
    "minimum": lax.min,
}
standard_op = lambda op: lambda *args, instr: op(*args)
ops = {name: standard_op(op) for name, op in standard_ops.items()}
# ops["tuple"] = lambda *args, instr: core.pack(args)
ops["convert"] = lambda x, instr: lax.convert_element_type(x, xla_client.XLA_ELEMENT_TYPE_TO_DTYPE[instr.shape.element_type])
ops["reshape"] = lambda x, instr: lax.reshape(x, instr.shape.dimensions)
ops["broadcast"] = lambda x, instr: lax.broadcast_in_dim(x, instr.shape.dimensions, instr.dimensions)
ops["concatenate"] = lambda *operands, **kw: lax.concatenate(operands, kw['instr'].dimensions[0])
ops["slice"] = lambda x, instr: lax.slice(x, *_slice_indices(instr))
ops["iota"] = lambda instr: lax.broadcasted_iota(xla_client.XLA_ELEMENT_TYPE_TO_DTYPE[instr.shape.element_type],
                                                 instr.shape.dimensions, instr.dimensions[0])
ops["compare"] = lambda x, y, instr: getattr(lax, instr.comparison_direction.lower())(x, y)
ops["reduce"] = lambda operand, init_val, computation, instr: lax.reduce(operand, init_val, computation, instr.dimensions)
ops["reverse"] = lambda x, instr: lax.rev(x, instr.dimensions)
ops["while"] = debug

def get_op(instr):
  try:
    return ops[instr.opcode]
  except KeyError:
    print(instr)
    raise

def check_shapes(val, shape):
  if shape.tuple_shapes:
    assert isinstance(val, core.JaxTuple)
    _map(check_shapes, val, shape.tuple_shapes)
  else:
    assert onp.shape(val) == tuple(shape.dimensions)
    assert onp.result_type(val) == xla_client.XLA_ELEMENT_TYPE_TO_DTYPE[shape.element_type]

def make_subcomp(hlo_computation):
  return partial(eval_hlo, DummyHloModule([hlo_computation]))
DummyHloModule = namedtuple("DummyHloModule", ["computations"])

def eval_hlo(hlo_proto, *args):
  sub_computations = hlo_proto.computations[:-1]
  computation = hlo_proto.computations[-1]

  def read(var):
    return env[var]

  def write(var, val):
    env[var] = val

  env = {}
  _map(write, parameters(computation), args)
  _map(write, constvars(computation), constvals(computation))
  for instruction in eqns(computation):
    in_vals = _map(read, instruction.operand_ids)
    if instruction.called_computation_ids:
      subcomp = get_computation(sub_computations, instruction.called_computation_ids[0])
      out_val = get_op(instruction)(*in_vals, instr=instruction,
                                    computation=make_subcomp(subcomp))
    else:
      out_val = get_op(instruction)(*in_vals, instr=instruction)
    check_shapes(out_val, instruction.shape)
    write(instruction.id, out_val)
  return read(computation.root_id)

[jkr26]

thanks! this is more than enough for me. Perfect!