Mojo function
scatter_nd_shape
scatter_nd_shape[input_type: DType, indices_type: DType](input: TileTensor[input_type, input.LayoutType, input.origin, address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size], updates: TileTensor[input_type, updates.LayoutType, updates.origin, address_space=updates.address_space, linear_idx_type=updates.linear_idx_type, element_size=updates.element_size], indices: TileTensor[indices_type, indices.LayoutType, indices.origin, address_space=indices.address_space, linear_idx_type=indices.linear_idx_type, element_size=indices.element_size]) -> IndexList[TileTensor[input_type, input.LayoutType, input.origin, address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size].rank]
Compute the output shape of a scatter_nd operation, and assert the inputs are compatible.
Parameters:
Args:
- input (
TileTensor): The input tensor. - updates (
TileTensor): The input tensor. - indices (
TileTensor): The indices tensor.
Returns:
IndexList: The output shape.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!