# Discretize only a certain arrrays in a tensor with TensorFlow

## Issue

I have the following array:-

``````import numpy as np
import tensorflow as tf
input = np.array([[-1.5, 1.0, 3.4, .5], [0.0, 3.0, 1.3, 0.0]])
layer = tf.keras.layers.Discretization(num_bins=2, epsilon=0.01)
layer(input)

<tf.Tensor: shape=(2, 4), dtype=int64, numpy=
array([[0, 1, 1, 1],
[0, 1, 1, 0]])>
``````

This discretizes the whole tensor. I would like to know if there is a way through which I can just discretize the second array in the tensor.

## Solution

We can create a mask based on the index of the array that needs to be discretized:

``````def get_mask(x, array_index):
x = tf.Variable(tf.ones_like(input, dtype=tf.float32))
indices = tf.Variable(array_index, dtype=tf.int32)
updates = tf.Variable(tf.zeros( (indices.shape[0], x.shape[1])), dtype=tf.float32)
``````

And calling

``````> mask = get_mask(input, np.array([[1]])) #second array
>
array([[1., 1., 1., 1.],
[0., 0., 0., 0.]])
``````

Then we can apply mask: `tf.cast(layer(input), tf.float32) * (1-mask) + input*mask` which returns:

``````array([[-1.5,  1. ,  3.4,  0.5],
[ 0. ,  1. ,  1. ,  0. ]]
``````

The above should work for any array and any array index to discretize.