HloGenerator

Generates StableHLO MLIR from any Model by:

  1. Creating a tape-recording execution context

  2. Running the model forward pass to record operations

  3. Converting the execution tape to a ComputeGraph

  4. Compiling the graph to StableHLO MLIR

Functions

Link copied to clipboard
suspend fun <D : DType, V> generate(model: Model<D, V, Tensor<D, V>, Tensor<D, V>>, sampleInput: Tensor<D, V>, functionName: String = "main"): StableHloModule

Generate StableHLO from any Model and a sample input tensor.

Link copied to clipboard
fun <D : DType, V> generateBlocking(model: Model<D, V, Tensor<D, V>, Tensor<D, V>>, sampleInput: Tensor<D, V>, functionName: String = "main"): StableHloModule

Blocking variant of generate for Java callers who cannot use suspend.