You could try something like this:
import tensorflow as tf
x = tf.constant([[
[ 1,2,3,5],
[ 2,4,6,3],
[ 0,7,1,2]],
[
[ 9,2,3,5],
[ 2,2,2,3],
[ 0,1,5,2]]
])
inputs = tf.keras.layers.Input((3, 4))
argmax = tf.argmax(inputs, axis=-1)
ij = tf.stack(tf.meshgrid(
tf.range(tf.shape(inputs)[0], dtype=tf.int64),
tf.range(tf.shape(inputs)[1], dtype=tf.int64),
indexing='ij'), axis=-1)
gather_indices = tf.concat([ij, tf.expand_dims(argmax, axis=-1)], axis=-1)
onehot = tf.tensor_scatter_nd_update(tf.zeros_like(x, dtype=tf.float32), gather_indices, tf.ones_like(argmax, dtype=tf.float32))
outputs = tf.keras.layers.Concatenate()([inputs,onehot])
model = tf.keras.Model(inputs, outputs)
print(model(x))
tf.Tensor(
[[[1. 2. 3. 5. 0. 0. 0. 1.]
[2. 4. 6. 3. 0. 0. 1. 0.]
[0. 7. 1. 2. 0. 1. 0. 0.]]
[[9. 2. 3. 5. 1. 0. 0. 0.]
[2. 2. 2. 3. 0. 0. 0. 1.]
[0. 1. 5. 2. 0. 0. 1. 0.]]], shape=(2, 3, 8), dtype=float32)
An explanation to what tf.meshgrid is and why it is needed can be found here. You can consider wrapping these operations in a custom Keras layer.
If you would use tf.reduce_max instead of tf.argmax, it would be way simpler:
import tensorflow as tf
x = tf.constant([[
[ 1,2,3,5],
[ 2,4,6,3],
[ 0,7,1,2]],
[
[ 9,2,3,5],
[ 2,2,2,3],
[ 0,1,5,2]]
])
inputs = tf.keras.layers.Input((3, 4))
onehot = tf.where(tf.not_equal(inputs, tf.reduce_max(inputs, keepdims=True, axis=-1)), tf.zeros_like(inputs), tf.ones_like(inputs))
outputs = tf.keras.layers.Concatenate()([inputs,onehot])
model = tf.keras.Model(inputs, outputs)
print(model(x))
tf.Tensor(
[[[1. 2. 3. 5. 0. 0. 0. 1.]
[2. 4. 6. 3. 0. 0. 1. 0.]
[0. 7. 1. 2. 0. 1. 0. 0.]]
[[9. 2. 3. 5. 1. 0. 0. 0.]
[2. 2. 2. 3. 0. 0. 0. 1.]
[0. 1. 5. 2. 0. 0. 1. 0.]]], shape=(2, 3, 8), dtype=float32)
But I guess there is a reason you want to use argmax.