Compilation with XLA can greatly improve the performance of your programs, but the TensorFlow interop has a number of known sharp corners.
tf.Variable
on a different deviceError message: INVALID_ARGUMENT: Trying to access resource (defined @ ) located in device CPU:0 from device GPU:0`
XLA cluster runs on exactly one device, and it can not read or write to tf.Variable
located on a different device. Usually this error message indicates that the variable was not placed on the right device to begin with. The error message should precisely specify the location of the offending variable.
NOTE: tf.Variable
of type int32
are always placed on a host, and can not be placed on a GPU. As a workaround, int64
can be used.
Error message: Support for TensorList crossing the XLA/TF boundary is not implemented
.
XLA supports tf.TensorArray
. However, the interconversion between TF and XLA representations is not implemented yet. This error often arises when the TensorArray
is used inside the compiled block, but the derivative is taken outside.
Workaround: compile the outermost scope which is taking the derivative.
Error message: XLA compilation requires a fixed tensor list size. Set the max number of elements. This could also happen if you're using a TensorArray in a while loop that does not have its maximum_iteration set, you can fix this by setting maximum_iteration to a suitable value
.
TF while loops created using tf.while_loop
support backpropagation by accumulating all intermediate results in a TensorArray
, but XLA only supports bounded TensorArray
s.
Workaround: all compiled while loops need to either have maximum_iterations
parameter set to a constant value known at compile time, or backpropagation disabled using back_prop=False
.
tf.TensorArray
is not supportedWrites into tf.TensorArray(..., dynamic_size=True)
are not compilable with XLA, as such writes require an unknown number of reallocations when the array exceeds the original bound.
Workaround: provide a statically known bound to your arrays.
XLA currently ignores TF seeds to random operations. This affects stateful TF random operations, such as tf.random.normal
, or tf.nn.dropout
. XLA will behave as if the compilation was seeded with a new unique seed at each run. This limitation does not apply to stateless random ops.