Searched refs:flattened_size (Results 1 – 2 of 2) sorted by relevance
47 int64 flattened_size = xla::Product(iota_shape.dimensions()); in Compile() local48 xla::XlaOp reshaped_condition = xla::Reshape(condition, {flattened_size}); in Compile()65 xla::XlaOp reshaped = xla::Reshape(iota, {flattened_size}); in Compile()77 to_concat.push_back(xla::Reshape(index_single_dim, {flattened_size, 1})); in Compile()
836 int32 flattened_size = 1; in ConvertSoftmaxOperator() local838 flattened_size *= input_shape.dims(i); in ConvertSoftmaxOperator()841 flattened_size, input_shape.dims(input_shape.dimensions_count() - 1)}; in ConvertSoftmaxOperator()878 int32 flattened_size = 1; in ConvertLogSoftmaxOperator() local880 flattened_size *= input_shape.dims(i); in ConvertLogSoftmaxOperator()883 flattened_size, input_shape.dims(input_shape.dimensions_count() - 1)}; in ConvertLogSoftmaxOperator()