[Doc] Create Colab tutorial that covers device placement APIs

[regrettable-username]

There are currently no tutorials covering device placement APIs such as:

• withDevice(_:_:perform:)
• withDefaultDevice(perform:)
• enableGPU()
• enableCPU()

[rxwei]

Please self-assign this once you start working on it. Thanks!

[regrettable-username]

I would, but I'm not able to assign the issue for some reason. I've started on the task and actually have a few questions:

1. Calls to enableCPU(), enableGPU() and enableTPU() seem to have no effect. Am I missing something? I'm assuming there is a scenario in which they are necessary.
2. withDevice(...) and withDefaultDevice(...) work just fine with .cpu and .gpu however, I was encountering a fatal error when trying to use .tpu on a TPU-backed colab notebook. I'll post the code and output below.
3. I was curious if there was a way to change the "default device"? It seems to be GPU on CUDA enabled toolchain (as expected).

The TPU issue I mentioned:

import TensorFlow
let a = Tensor<Float>(randomNormal: [1000, 1000])
let b = Tensor<Float>(randomNormal: [1000, 1000])
enableTPU()

withDevice(.tpu) {
  let c = a * b
//   print(c.mean())
}

Output:

Fatal error: Unexpectedly found nil while unwrapping an Optional value
Current stack trace:
0    libswiftCore.so                    0x00007fcae2e180c0 _swift_stdlib_reportFatalError + 69
1    libswiftCore.so                    0x00007fcae2d292f6 <unavailable> + 2777846
2    libswiftCore.so                    0x00007fcae2d29675 <unavailable> + 2778741
3    libswiftCore.so                    0x00007fcae2ba7480 _fatalErrorMessage(_:_:file:line:flags:) + 19
4    libswiftTensorFlow.so              0x00007fcae023dd80 _ExecutionContext.currentDeviceName.getter + 1761
5    libswiftTensorFlow.so              0x00007fcae025d0c0 _TFCOpSetDeviceFromScope(_:_:) + 156
6    libswiftTensorFlow.so              0x00007fcae025d0b0 _swift_tfc_OpSetDeviceFromScope + 9
Current stack trace:
	frame #7: 0x00007fcb0b62c2e1 $__lldb_expr28`specialized thunk for @callee_guaranteed () -> (@error @owned Error) [inlined] closure #1 () -> () in __lldb_expr_27 at <Cell 1>:7
	frame #10: 0x00007fcb0b62c294 $__lldb_expr28`main at <Cell 1>:6:1

[rxwei]

This tutorial doesn't need to work with TPU. It's not a publicly available feature yet.

[rxwei]

1. Calls to enableCPU(), enableGPU() and enableTPU() seem to have no effect. Am I missing something? I'm assuming there is a scenario in which they are necessary.

I think you need to change the runtime type from CPU to GPU (Runtime -> Change runtime type). Otherwise things will fall back to CPU.

3. I was curious if there was a way to change the "default device"? It seems to be GPU on CUDA enabled toolchain (as expected).

There's no way to change it, which is kind of the point for "default". When you need to change the device, use withDevice(_:perform:) or enableGPU() and a new device kind will be pushed to a thread-local stack in the runtime that remembers scope-specific device configuration.

[regrettable-username]

Ok, that makes sense. Thanks for clarifying. I'll have a PR shortly.

[BradLarson]

This is an older issue that has a matching (closed) PR, so I'll close this in favor of the PR that came from it. If this is still useful to have, we can reopen the corresponding PR.