from a conversation with Daniel after https://sflean.group/

  • pytorch -> triton -> tritonIR -> ptx -> sass
    • or eager-mode: pytorch -> prebuilt cuda kernels
  • jax -> jaxpr -> xla -> llvm ir -> ptx -> sass
    • can’t have if statements in jax - branches are ?collapsed? at compiletime
    • there is a keyword in jax that gets around this
    • it lowers the branch to the corresponding function in xla