View source on GitHub |
Select values from x
at the 1-D indices
along the given axis.
tf.keras.ops.take_along_axis(
x, indices, axis=None
)
Args | |
---|---|
x
|
Source tensor. |
indices
|
The indices of the values to extract. |
axis
|
The axis over which to select values. By default, the flattened input tensor is used. |
Returns | |
---|---|
The corresponding tensor of values. |