JAX Framework

Note

DALI keeps the buffers it exports through DLPack alive by synchronizing the device from a background thread. This invalidates the CUDA graphs that XLA captures for its command buffers, which surfaces as CUDA_ERROR_STREAM_CAPTURE_INVALIDATED.

The examples below disable command buffers with the XLA_FLAGS environment variable. This is necessary when combining DALI with jax.jit.

• Getting started with JAX and DALI
  • Prerequisites
  • Quick start
  • Running DALI with JAX
  • Batching
  • GPU acceleration
  • Hardware accelerated decoding
  • Random shuffling
  • Augmentations
  • Multiple GPUs
  • Technical details
  • Where to go next
• Training a neural network with DALI and JAX
• Training with multiple GPUs
  • Training with automatic parallelization
  • Training with pmapped iterator
• Training neural network with DALI and Flax
• Multiple GPUs with DALI and FLAX
• Training neural network with DALI and Paxml
• Getting started with T5X and DALI
  • Prerequisites
  • Peekable iterator for T5X
  • Where to go next