Searched refs:torch_dispatch_class (Results 1 – 3 of 3) sorted by relevance
| /external/pytorch/torch/_library/ |
| D | simple_registry.py | 65 self, torch_dispatch_class: type, func: Callable
67 if self.find(torch_dispatch_class):
71 self._data[torch_dispatch_class] = func
74 del self._data[torch_dispatch_class]
78 def find(self, torch_dispatch_class): argument
79 return self._data.get(torch_dispatch_class, None)
82 def find_torch_dispatch_rule(op, torch_dispatch_class: type) -> Optional[Callable]:
84 torch_dispatch_class
|
| D | custom_ops.py | 462 self, torch_dispatch_class: Any, fn: Optional[Callable] = None, /
474 if torch_dispatch_class not in self._torch_dispatch_fns:
477 return self._torch_dispatch_fns[torch_dispatch_class](
482 self._name, torch_dispatch_class, inner
484 self._torch_dispatch_fns[torch_dispatch_class] = fn
|
| /external/pytorch/torch/ |
| D | library.py | 187 def _register_torch_dispatch_rule(self, op_name, torch_dispatch_class, fn): argument
205 handle = entry.torch_dispatch_rules.register(torch_dispatch_class, fn)
933 torch_dispatch_class: Any,
993 return opdef.register_torch_dispatch(torch_dispatch_class, func)
1003 use_lib._register_torch_dispatch_rule(op_name, torch_dispatch_class, func)
|