матрица умножения тензорного потока внутри заданной оси

Question

У меня есть тензор формы (32,1024,40,3) и вектор шаговой функции размером 40 (т.е. [1,1,1,...,1,0,0,...,0]), и я хотел бы умножить матрицу в пределах заданной оси.
Это мой код с ошибкой'Tensor' object does not support item assignment в строке 716:

def weight_knn(dist_from_point, weight_param):
    given_axis = dist_from_point.shape[-2]
    step = np.ones(given_axis)
    step[weight_param:] = 0
    dim_array = tf.reshape(tf.ones((1, tf.rank(dist_from_point)),tf.uint16),[-1])
    dim_array[-2] = -1
    step_reshaped = step.reshape(dim_array)
    return dist_from_point*step_reshaped

Где идея умножения вдоль определенной оси была взята из этого фрагмента с фрагментами
И цель состоит в том, чтобы иметь возможность умножитьфункция шага с мультиизмеренным тензором в заданном dim