mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
c9023e3b12
Summary:
With this PR, we can now support left and right shift operators in the JIT engine for <int, int> and <Tensor, int>.
Updated tests pass as expected:
```
> python test/test_jit.py
...
Ran 2427 tests in 84.861s
OK (skipped=139, expected failures=1)
```
Running the following code with Python results in the output below:
```
> cat ~/expressions.py
import torch
torch.jit.script
def fn(a, b):
# type: (int, int)
return (
a << b, # supported
b >> a, # supported
a & b,
a | b,
a ^ b
)
print(fn.graph)
```
```
> python ~/expressions.py
graph(%a.1 : int,
%b.1 : int):
%4 : int = aten::leftshift(%a.1, %b.1) # /home/ince/expressions.py:7:8
%7 : int = aten::rightshift(%b.1, %a.1) # /home/ince/expressions.py:8:8
%10 : int = aten::__and__(%a.1, %b.1) # /home/ince/expressions.py:9:8
%13 : int = aten::__or__(%a.1, %b.1) # /home/ince/expressions.py:10:8
%16 : int = aten::__xor__(%a.1, %b.1) # /home/ince/expressions.py:11:8
%17 : (int, int, int, int, int) = prim::TupleConstruct(%4, %7, %10, %13, %16)
return (%17)
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/34563
Differential Revision: D20434209
Pulled By: tugrulince
fbshipit-source-id: 886386c59755106e17b84778b8e495b80a6269cd
|
||
|---|---|---|
| .. | ||
| builtin_functions.cpp | ||
| builtin_functions.h | ||
| canonicalize_modified_loop.cpp | ||
| canonicalize_modified_loop.h | ||
| code_template.h | ||
| concrete_module_type.cpp | ||
| concrete_module_type.h | ||
| convert_to_ssa.cpp | ||
| convert_to_ssa.h | ||
| edit_distance.cpp | ||
| edit_distance.h | ||
| error_report.cpp | ||
| error_report.h | ||
| exit_transforms.cpp | ||
| exit_transforms.h | ||
| function_schema_parser.cpp | ||
| function_schema_parser.h | ||
| inline_loop_condition.cpp | ||
| inline_loop_condition.h | ||
| ir_emitter.cpp | ||
| ir_emitter.h | ||
| lexer.cpp | ||
| lexer.h | ||
| mini_environment.h | ||
| parse_string_literal.h | ||
| parser_constants.h | ||
| parser.cpp | ||
| parser.h | ||
| resolver.h | ||
| schema_matching.cpp | ||
| schema_matching.h | ||
| schema_type_parser.cpp | ||
| schema_type_parser.h | ||
| script_type_parser.cpp | ||
| script_type_parser.h | ||
| source_range.cpp | ||
| source_range.h | ||
| string_to_type.cpp | ||
| strtod.cpp | ||
| strtod.h | ||
| sugared_value.cpp | ||
| sugared_value.h | ||
| tracer.cpp | ||
| tracer.h | ||
| tree_views.h | ||
| tree.h | ||