XLA dispatcher

Looping primitives

While

  • The shape of the init parameter must be statically determined and remain fixed during iterations.
  • Result is initialized with the init parameter, and updated at each iteration of the body. If condition not met at the first iteration, returns init.

Map

Shapes

There does not seem to be dynamic shapes planned in XLA. But this in JAX. Must be a special kind of dynamic shape, i.e. that can be determined at tracing time.

Links to this note