EXLA

View Source EXLA.Computation (EXLA v0.4.1)

Wrapper around XLA's computation.

Link to this section Summary

Functions

compile(computation, client, argument_shapes, options \\ [])

Compiles a computation into an executable.

Link to this section Functions

Link to this function

compile(computation, client, argument_shapes, options \\ [])

View Source

Compiles a computation into an executable.

options

 Options

  • :device_id - the device id to compile to and run the executable on. Defaults to the :default_device_id on the client. If :num_replicas or :num_partitions are given, this option is ignored and the device id is set to -1.

  • :num_replicas - the number of replicas this computation will run on. It defaults to 1 but you can set it if you want to enable single-program multiple data.

  • :use_spmd - enables Single-Program Multiple-Data partioning. This is set to true if :num_partitions is more than one, otherwise is false.

Currently those options do not have an effect as they related to running the same compiled executable on multiple replicas.

Some options apply to TPU only:

  • :num_partitions - the number of partitions this computation will run on.