Lines Matching refs:dispatch_key

170             dispatch_key=DispatchKey.Undefined,
184             dispatch_key=k,
327         return [backend.dispatch_key for backend in backends] + [
343         return backend_index.dispatch_key
363         dispatch_key = get_static_dispatch_backend(f, index)
364         if dispatch_key is not None:
374         for dispatch_key in static_dispatch_keys(backends)
487             and b.dispatch_key in STRUCTURED_DISPATCH_KEYS
956         dispatch_key = "c10::computeDispatchKey(dtype, layout, device)"
1442         for dispatch_key, backend_idx in backend_indices.items():
1446                 dispatch_keys.add(dispatch_key)
1533     dispatch_key: DispatchKey,
1621                     "dispatch_namespace": dispatch_key.lower(),
1635     dispatch_key: DispatchKey,
1655             "native", dispatch_key.lower()
1806     for dispatch_key in dispatch_keys:
1807         fm = cuda_fm if is_cuda_dispatch_key(dispatch_key) else cpu_fm
1808         if dispatch_key in functions_keys:
1815                     "dispatch_key": str(dispatch_key),
1824                     "dispatch_namespace": dispatch_key.lower(),
1827                         dispatch_key=dispatch_key,
1828                         backend_idx=backend_indices[dispatch_key],
1958     for dispatch_key in dispatch_keys:
1959         if dispatch_key not in functions_keys:
1962         dispatch_namespace = dispatch_key.lower()
1970                         backend_indices[dispatch_key],
1995         fm = cuda_fm if is_cuda_dispatch_key(dispatch_key) else cpu_fm
2002                 "dispatch_key": str(dispatch_key),
2205     for dispatch_key in dispatch_keys:
2206         fm = cuda_fm if is_cuda_dispatch_key(dispatch_key) else cpu_fm
2229                     elif g.structured and dispatch_key in (
2239                         dispatch_key
2243                     if dispatch_key in functions_keys:
2254                 if dispatch_key == DispatchKey.CompositeExplicitAutogradNonFunctional:
2256                 if dispatch_key in functions_keys:
2260         backend_index = backend_indices[dispatch_key]
2270         dispatch_namespace = str(dispatch_key).lower()
2275             dispatch_key != DispatchKey.CompositeImplicitAutogradNestedTensor
2281             dispatch_key=dispatch_key,
2294                 if is_cuda_dispatch_key(dispatch_key)
2307             if not g.out.ufunc_inner_loop or not is_ufunc_dispatch_key(dispatch_key):
2310             if dispatch_key is DispatchKey.CPU:
2318                             g, backend_indices[dispatch_key]
2331             elif dispatch_key is DispatchKey.CUDA:
2343                             g, backend_indices[dispatch_key]
2358         if dispatch_key in (DispatchKey.CPU, DispatchKey.CUDA):
2373                 dispatch_key,
2418                         func, structured_func_group_dict, dispatch_key, backend_indices
2425                 extra_cuda_headers if is_cuda_dispatch_key(dispatch_key) else ""
2433                     dispatch_key,