# Distributed Fast Fourier Transform in TensorFlow

DevFeed: [Distributed Fast Fourier Transform in TensorFlow](<https://devfeed.tech/articles/distributed-fast-fourier-transform-in-tensorflow-7380.md>)

Original publisher: [Read original article](<https://blog.tensorflow.org/2023/08/distributed-fast-fourier-transform-in-tensorflow.html>)

Author: TensorFlow Blog (noreply@blogger.com)

Published: 2023-08-24T17:00:00Z

Content type: tutorial

Language: en

Sources: [The TensorFlow Blog](<https://devfeed.tech/sources/the-tensorflow-blog.md>)

Topics: [AI, ML & Data Engineering](<https://devfeed.tech/topics/ai-ml-data-engineering.md>), [GPU](<https://devfeed.tech/topics/gpu.md>), [NCCL](<https://devfeed.tech/topics/nccl.md>)

Tags: [api](<https://devfeed.tech/tags/api.md>), [gpu](<https://devfeed.tech/tags/gpu.md>), [machine-learning](<https://devfeed.tech/tags/machine-learning.md>), [nccl](<https://devfeed.tech/tags/nccl.md>), [performance](<https://devfeed.tech/tags/performance.md>), [tensorflow](<https://devfeed.tech/tags/tensorflow.md>), [tensorflow-core](<https://devfeed.tech/tags/tensorflow-core.md>)

## AI overview

TensorFlow v2 adds native Distributed FFT support through DTensor. The article explains how sharded tensors can use existing FFT operations across multiple devices, and reports that communication and data shuffling dominate runtime in the demonstrated setup.

## Source excerpt

Posted by Ruijiao Sun, Google Intern - DTensor team Fast Fourier Transform is an important method of signal processing, which is commonly used in a number of ways, including speeding up convolutions, extracting features, and regularizing models. Distributed Fast Fourier Transform (Distributed FFT) offers a way to compute Fourier Transforms in models that work with image-like datasets that are too large to fit into the memory of a single accelerator device. In a previous Google Research Paper, "Large-Scale Discrete Fourier Transform on TPUs" by Tianjian Lu, a Distributed FFT algorithm was implemented for TensorFlow v1 as a library. This work presents the newly added native support in TensorFlow v2 for Distributed FFT, through the new TensorFlow distribution API, DTensor. About DTensor DTensor is an extension to TensorFlow for synchronous distributed computing. It distributes the program and tensors through a procedure called Single program, multiple data (SPMD) extension. DTensor offers an uniform API for traditional data and model parallelism patterns used widely in Machine Learning. Example Usage The API interface for distributed FFT is the same as the original FFT in TensorFlow. Users just need to pass a sharded tensor as an input to the existing FFT ops in TensorFlow, such as tf.signal.fft2d. The output of a distributed FFT becomes sharded too. import TensorFlow as tf from TensorFlow.experimental import dtensor # Set up devices device_type = dtensor.preferred_device_type() if device_type == 'CPU': cpu = tf.config.list_physical_devices(device_type) tf.config.set_logical_device_configuration(cpu[0], [tf.config.LogicalDeviceConfiguration()] * 8) if device_type == 'GPU': gpu = tf.config.list_physical_devices(device_type) tf.config.set_logical_device_configuration(gpu[0], [tf.config.LogicalDeviceConfiguration(memory_limit=1000)] * 8) dtensor.initialize_accelerator_system() # Create a mesh mesh = dtensor.create_distributed_mesh(mesh_dims=[('x', 1), ('y', 2), ('z', 4)],