tf.function wraps a TensorFlow function so TensorFlow can compile it into a more efficient computation. The default mode is graph mode; XLA compilation is optional. section 3.3.1
Compile a dense operation
```r dense <- tf_function(\(inputs, W, b) { tf$nn$relu(tf$matmul(inputs, W) + b) }) ``` The first call takes longer because TensorFlow compiles the function. Later calls reuse the compiled program and run faster. section 3.3.1
Use tf.function with these practical expectations:
- Debug the ordinary TensorFlow function first; compilation changes the executed code and can make debugging more difficult. section 3.3.1
- Choose the default graph mode for compilation, or set jit_compile = TRUE to request XLA, which may run faster but takes longer to compile initially. section 3.3.1
- Wrap gradient-based steps as well; a GradientTape can compute the loss and gradients inside a compiled stochastic gradient-ascent step. section 10.2.1
dense <- tf_function(\(inputs, W, b) {
tf$nn$relu(tf$matmul(inputs, W) + b)
})
Runnable R syntax for wrapping a TensorFlow dense function with tf_function. section 3.3.1
dense <- tf_function(jit_compile = TRUE, \(inputs, W, b) {
tf$nn$relu(tf$matmul(inputs, W) + b)
})
Runnable R syntax for requesting XLA compilation with jit_compile = TRUE. section 3.3.1
