Searched refs:has_batch (Results 1 – 4 of 4) sorted by relevance
93 const bool has_batch = data->shape().rank() == 2; in Run() local117 if (has_batch && slice->slice_limits(batch_dim) != in Run()141 has_batch ? sort->operand(0)->shape().dimensions(batch_dim) : 1; in Run()144 if (has_batch && sort_dim == 0) { in Run()151 has_batch ? ShapeUtil::MakeTupleShape( in Run()167 if (has_batch && sort_dim == 0) { in Run()
63 bool has_batch; member
69 if (stg->has_batch) in init_perf_monitor()432 stg->has_batch = true; in st_InitPerfMonitorGroups()
2260 const bool has_batch = hlo->shape().tuple_shapes(0).dimensions_size() == 2; in HandleTopK() local2283 {b_.getInt64(has_batch ? input->shape().dimensions(0) : 1), in HandleTopK()