pytorch
6a2838ee - [jit] jit._drop fun modifier to allow in jit class non-jit decl funs (#93012)

Commit
3 years ago
[jit] jit._drop fun modifier to allow in jit class non-jit decl funs (#93012) `@torch.jit.unused` and `@torch.jit.ignore` do not allow to keep in torch scripted class member function, that has non scriptable declaration (e.g. return type) Adding FunctionModifier _DROP to allow fully skip those functions from scripting and keep them in the code of the scripted class. E.g. it can be used for: ``` @torch.jit._drop def __fx_create_arg__(self, tracer: torch.fx.Tracer) -> torch.fx.node.Argument: # torch.fx classes are not scriptable return tracer.create_node( "call_function", CFX, args=(tracer.create_arg(self.features),), kwargs={}, ) def __iter__(self) -> Iterator[torch.Tensor]: return iter(self.a) ``` Testing: Added test case in `test/jit/test_types.py` with non-scriptable type annotations (fx.* classes) that fails before fix and passes after. ``` python test/test_jit.py ``` Differential Revision: [D42774830](https://our.internmc.facebook.com/intern/diff/D42774830) Pull Request resolved: https://github.com/pytorch/pytorch/pull/93012 Approved by: https://github.com/davidberard98
Author
Committer
Parents
Loading