How to assign one of the outputs of a layer to be 0 - tensorflow

Question

I have model built in keras, tensorflow 2.0, I want to force the first output of its last layer to be zero. How can I do that?
Here is the model I built and my try for that process:
import tensorflow as tf
X = Input(shape=(32,))
Y = Dense(256, activation= 'relu', kernel_initializer=ini)(X)
Here is what I tried to do. I tried to create another layer whose all values are zeros except the first values which is similar to the output of the layer Y
def zeros_assign(Y):
a = tf.zeros(256,)
aa = aa[0].assign[-Y[0]]
out = aa + Y
return out
Y_out = Lambda(lambda x: channel(x))(Y)
But when I run that layer, it give an error as following:
TypeError: 'Functional' object does not support item assignment

Answer 1

Following is a solution where the BlockFirstNode layer pass all the output from the previous layer, except the first node which is blocked (set to 0). This is as suggested by Dr. Snoopy above by multiplying by zero.
Code example:
import tensorflow as tf
import numpy as np
class BlockFirstNode(tf.keras.layers.Layer):
def __init__(self) -> None:
super().__init__()
def build(self, input_shape):
output_len = input_shape[1]
bloc_vec_np = np.array([0]+[1]*(output_len-1)).reshape(1, output_len)
self.block_vec = tf.convert_to_tensor(bloc_vec_np, dtype=tf.float32)
def call(self, inputs):
return inputs*self.block_vec
X = tf.keras.layers.Input(shape=(32,))
Y = tf.keras.layers.Dense(5, activation= 'sigmoid')(X)
block_output = BlockFirstNode()(Y)
model = tf.keras.models.Model(inputs=X, outputs=[Y, block_output])
pred_y, pred_block = model.predict(np.ones((1,32)))
print(f'Y: {np.round(pred_y,2)}')
# Y: [[0.68 0.59 0.88 0.08 0.48]]
print(f'block_output: {np.round(pred_block,2)}')
# block_output: [[0. 0.59 0.88 0.08 0.48]]
Here we can see that all output are the same as the layer before, except the first node that is forced to zero.
EDIT (2023-02-04):
Now work for any n number of blockings, as requested by Sajjad.
import tensorflow as tf
import numpy as np
class BlockNFirstNode(tf.keras.layers.Layer):
def __init__(self, n_blocking) -> None:
super().__init__()
self.n_blocking = n_blocking
def build(self, input_shape):
output_len = input_shape[1]
if self.n_blocking>output_len:
raise ValueError("n_blocking cannot be larger than layer size")
bloc_vec_list = [0]*self.n_blocking+[1]*(output_len-self.n_blocking)
bloc_vec_np = np.array(bloc_vec_list).reshape(1, output_len)
self.block_vec = tf.convert_to_tensor(bloc_vec_np, dtype=tf.float32)
def call(self, inputs):
return inputs*self.block_vec
X = tf.keras.layers.Input(shape=(32,))
Y = tf.keras.layers.Dense(5, activation= 'sigmoid')(X)
block_output = BlockNFirstNode(n_blocking=2)(Y)
model = tf.keras.models.Model(inputs=X, outputs=[Y, block_output])
pred_y, pred_block = model.predict(np.ones((1,32)))
print(f'Y: {np.round(pred_y,2)}')
# Y: [[0.68, 0.59, 0.88, 0.08, 0.48]]
print(f'block_output: {np.round(pred_block,2)}')
# block_output: [[0, 0, 0.88, 0.08, 0.48]]