Device Placement and Dtypes with jax.device_put

1. Discovering Your Computational Environment

Before you can direct traffic, you need to see the roads. In JAX, this means querying the available hardware. JAX provides a straightforward API to inspect the CPUs, GPUs, or TPUs it can access.

The primary functions for this are:

• jax.devices(): Returns a list of all available devices. In a multi-host environment (like a TPU pod), this shows devices across all hosts.
• jax.local_devices(): Returns a list of devices attached to the current host/process. For single-machine setups, this is often identical to jax.devices().
• jax.device_count() and jax.local_device_count(): Return the number of devices from the functions above.
• jax.default_backend(): Returns the name of the default platform JAX has chosen (e.g., 'gpu', 'tpu', 'cpu').

The following reading provides a clear guide on using these functions.

Let's see this in action. You can run this code in your own environment to see what JAX detects.

import jax

print(f"JAX default backend: {jax.default_backend()}")
print(f"All available JAX devices: {jax.devices()}")
print(f"Number of local devices: {jax.local_device_count()}")




# You can also get devices for a specific backend
try:
    print(f"GPU devices: {jax.devices('gpu')}")
except RuntimeError:
    print("No GPU backend found.")

try:
    print(f"CPU devices: {jax.devices('cpu')}")
except RuntimeError:
    print("No CPU backend found.")

Each item in the list returned by jax.devices() is a Device object, which you can use to direct your computations.

2. Explicit Placement and Inspection

As we discussed, JAX automatically places newly created arrays on a default device, typically the first available accelerator. While convenient, for reproducible performance and complex workflows, explicit control is essential. This is the role of jax.device_put().

jax.device_put(x, device=None) takes a host-side object (like a NumPy array or Python scalar) and returns a jax.Array residing on the specified device. If device is None, the default device is used. This is the primary mechanism for moving data from CPU memory (the "host") to accelerator memory (the "device").

Once you have a jax.Array, you can inspect its properties using the .device() method and .dtype attribute.

Let's synthesize these concepts. The following code demonstrates the full loop: query devices, create data on the host, move it to a specific device, and verify.

import jax
import jax.numpy as jnp
import numpy as np




# 1. Discover devices
devices = jax.local_devices()
print(f"Local devices: {devices}")

if not devices:
    print("No JAX devices found. Exiting.")
else:



    # Use the first available device as our target
    target_device = devices[0]
    print(f"Targeting device: {target_device}")




    # 2. Create data on the host (CPU)
    host_data = np.linspace(0, 1, 5).astype(np.float32)
    print(f"\nCreated NumPy array (on host): {host_data}, type: {host_data.dtype}")




    # 3. Explicitly place data on the target device
    device_data = jax.device_put(host_data, device=target_device)




    # 4. Inspect the resulting jax.Array
    print(f"Placed JAX array: {device_data}")
    print(f"Array is on device: {device_data.device()}")
    print(f"Array dtype is: {device_data.dtype}")

This pattern of creating host-side data (e.g., from a data loader) and then using jax.device_put to transfer it to the accelerator is fundamental to JAX programming.

3. Data Types and Promotion Rules

The learning outcome also specifies inspecting dtypes. While arr.dtype is simple enough, the behavior of dtypes in JAX warrants a deeper look, especially given your mathematics background. Like other numerical libraries, JAX has rules for determining the resulting dtype when you perform operations on arrays of different types (e.g., adding a float16 to an int32).

JAX's type promotion system is designed to be predictable and to prevent silent precision loss where possible. It generally follows a "weak" type promotion for operations between Python scalars and JAX arrays, and a more standard promotion for operations between JAX arrays.

The rules can be visualized as a promotion lattice. When two types are combined, the result is their "join" in this lattice—the lowest common type that is "above" both.

The key takeaway from this graph is that integers combined with floats become floats, and lower-precision types are promoted to match higher-precision ones. For example, i* (integer types) and f* (float types) both lead to f*. Combining an f32 and a c64 will result in a c64.

A Practical "Gotcha"

A common source of unexpected behavior comes from Python scalars. By default, a literal like 0.1 is a 64-bit float in Python. If you combine it with a 32-bit JAX array, the array will be promoted to 64-bit, which can halve your performance on hardware optimized for 32-bit math.

x_f32 = jnp.ones(3, dtype=jnp.float32)




# Promotion to float64 due to Python scalar
y = x_f32 + 0.1
print(f"Adding a Python float literal: x.dtype={x_f32.dtype}, y.dtype={y.dtype}")




# To preserve float32, be explicit
y_f32 = x_f32 + jnp.float32(0.1)
print(f"Adding an explicit JAX float32: y.dtype={y_f32.dtype}")

For most deep learning work, you'll want to stay in 32-bit (or even 16-bit) precision. Being mindful of how you introduce scalars is key. If you genuinely need 64-bit precision, you can enable it globally with a config flag (jax.config.update("jax_enable_x64", True)), but the more idiomatic and safer approach is to manage dtypes explicitly.

4. The Host-Device Memory Model

Let's formalize the distinction between the "host" and the "device," as it's central to understanding performance.

• Host: The CPU that runs your main Python script. Standard Python objects and NumPy arrays reside in the host's RAM.
• Device: The accelerator (GPU/TPU) where high-performance computation happens. JAX arrays that are part of accelerated computations live in the device's dedicated memory (e.g., VRAM).

Every jax.device_put call initiates a data transfer across a physical bus (like PCIe) from host RAM to device VRAM. This transfer has a cost. The key to high performance in JAX is to minimize these transfers. The ideal workflow is:

1. Transfer In: Move your raw data (e.g., a batch of images or text) to the device once.
2. Compute on Device: Execute your entire computation graph (e.g., the forward and backward pass of a neural network) purely on the device using JAX's transformed functions.
3. Transfer Out: Only bring back the final, small results you need to the host (e.g., the scalar loss value for logging).

The video from our previous lesson provided a great benchmark showing the overhead of implicit transfers when a JAX operation is called on NumPy arrays.

Key Takeaways

In this lesson, you have gained the tools to precisely manage the location and type of your data on a single accelerator.

• Inspect Your Environment: Use jax.devices() and its variants to know what hardware is available to your JAX program.
• Inspect Your Arrays: Use the arr.device() method and arr.dtype attribute to verify where your arrays live and what their precision is.
• Control Placement: jax.device_put() is your primary tool for explicitly moving data from the host (CPU) to the accelerator (device).
• Control Precision: Be mindful of JAX's type promotion rules, especially when interacting with Python scalars, to maintain the desired precision for optimal performance.
• Minimize Transfers: A core principle for performance is to move data to the device, compute as much as possible there, and only transfer final results back to the host.

With this solid understanding of jax.Array objects and their device placement, we are now ready to start working with functions. In the next lesson, we will tackle a core challenge in transitioning to JAX: converting stateful, imperative Python code into the pure functions that JAX's transformation-based system requires.