How to apply a function (similar to map) on a tensor

What is a tensor?

Tensor is a fundamental mathematical concept commonly used in various fields like mathematics, physics, machine learning, deep learning, etc. They can be considered multidimensional arrays or generalized vectors that store and represent data of various types and dimensions. Tensors have the following features:

  • Have multiple dimensions

  • Composed of individual elements (like numbers, scalars, etc.)

  • Stores data of various datatypes

  • Supports various mathematical operations

Tensor manipulation is a fundamental aspect of deep learning and data processing, and being able to apply a function on a tensor is a common operation. This is similar to applying the map function of Python to each element of an iterable

This answer will explore two different techniques for applying a function to a tensor.

  • Using numpy library

  • Using tensorflow library

Using the numpy library

Here’s how we can apply a function to a tensor using numpy library. Run the code to see the output:

import numpy as np
# Create a NumPy array (tensor)
tensor = np.array([1, 2, 3, 4, 5])
# Define a function to apply 
def myFunction(x):
    return x + 2
# Use NumPy's vectorized operations to apply the function to each element
result = np.vectorize(myFunction)(tensor)
print(result)

Code explanation

  • Line 1: We import the numpy library and assign it an alias np. It is a popular library in Python for numerical computations.

  • Line 4: NumPy array (or tensor) is created using np.array function.

  • Line 11: This line is where the magic happens:

    • The np.vectorize function of NumPy takes a Python function (my_func in this case) and returns a new function that vectorizes the original function.

    • The vectorized function is immediately applied to the tensor created earlier. It means the vectorized function is applied element-wise on each element of the tensor.

    • The result of the vectorized function applied to the tensor is stored in result variable. The result will be a tensor where each element is the result of applying the vectorized function to the corresponding element in the original tensor.

Using the tensorflow library

The other way to apply a function on a tensor is to use tensorflow library:

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
# Create a TensorFlow tensor
tensor = tf.constant([1, 2, 3, 4, 5])
# Define a function to apply to each element of the tensor
def myFunction(x):
    return x ** 2
# Use TensorFlow's map_fn to apply the function to each element
result = tf.map_fn(myFunction, tensor)
# Convert the result to a NumPy array for printing
resultNumpy = result.numpy()
print(resultNumpy)

Code explanation

  • Line 1: We import the tensorflow library and assign it an alias tf. This allows us to use its functions and classes in our code.

  • Line 4: We create a TensorFlow tensor.

  • Line 11: We use TensorFlow tf.map_fn function to apply my_function to each element of tensor. The resulting tensor is stored in result variable.

  • Line 14: We convert the result tensor into a NumPy array since it is an easy way to print the resulting tensor.

In conclusion, applying a function to a tensor can be achieved in various ways, each suited to different scenarios. Choose the method that best suits your specific needs and the size of your tensors to ensure optimal performance.

Free AI Mock Interviews
Coding Interview
Coding PatternsFree Interview
Gain insights and practical experience with coding patterns through targeted MCQs and coding problems, designed to match and challenge your expertise level.
System Design
You TubeFree Interview
Learn to design a video streaming platform like YouTube by tackling functional and non-functional requirements, core components, and high-level to detailed design challenges.

Free Resources

Copyright ©2025 Educative, Inc. All rights reserved