Searched refs:cudnn_call (Results 1 – 2 of 2) sorted by relevance
441 const HloCustomCallInstruction* cudnn_call) { in GetGpuConvConfig() argument444 TF_ASSIGN_OR_RETURN(descriptor.kind, GetCudnnConvKind(cudnn_call)); in GetGpuConvConfig()446 cudnn_call->backend_config<CudnnConvBackendConfig>()); in GetGpuConvConfig()447 descriptor.operand0_shape = cudnn_call->operand(0)->shape(); in GetGpuConvConfig()448 descriptor.operand1_shape = cudnn_call->operand(1)->shape(); in GetGpuConvConfig()449 descriptor.result_shape = cudnn_call->shape().tuple_shapes(0); in GetGpuConvConfig()450 descriptor.scratch_size = cudnn_call->shape().tuple_shapes(1).dimensions(0); in GetGpuConvConfig()451 descriptor.window = cudnn_call->window(); in GetGpuConvConfig()452 descriptor.dnums = cudnn_call->convolution_dimension_numbers(); in GetGpuConvConfig()453 descriptor.feature_group_count = cudnn_call->feature_group_count(); in GetGpuConvConfig()[all …]
139 const HloCustomCallInstruction* cudnn_call);