(*args, **kwargs)
| 8 | try_caching=False): |
| 9 | @functools.wraps(orig_fn) |
| 10 | def wrapper(*args, **kwargs): |
| 11 | if not handle.is_active(): |
| 12 | return orig_fn(*args, **kwargs) |
| 13 | |
| 14 | input_types = [ |
| 15 | v.data.type() for v in list(args) + list(kwargs.values()) |
| 16 | if utils.is_fp_tensor(v) |
| 17 | ] |
| 18 | #print('wrapper: orig_fn:{}, input_types:{}'.format(orig_fn, input_types)) |
| 19 | input_type = input_types[0] |
| 20 | |
| 21 | if try_caching and handle.has_cache: |
| 22 | args = list(args) |
| 23 | for i in range(len(args)): |
| 24 | if utils.should_cache(args[i]): |
| 25 | args[i] = utils.cached_cast(cast_fn, args[i], handle.cache) |
| 26 | for k in kwargs: |
| 27 | if utils.should_cache(kwargs[k]): |
| 28 | kwargs[k] = utils.cached_cast(cast_fn, kwargs[k], handle.cache) |
| 29 | new_args = utils.casted_args(cast_fn, |
| 30 | args, |
| 31 | kwargs) |
| 32 | output = orig_fn(*new_args, **kwargs) |
| 33 | |
| 34 | #if output.type() != input_type: |
| 35 | # print('ori output type: {}, input type: {}'.format(output.type(), input_type)) |
| 36 | # return output.type(input_type) |
| 37 | #return output |
| 38 | return cast_output(output, input_type, verbose=False) |
| 39 | |
| 40 | return wrapper |
| 41 |
nothing calls this directly
no test coverage detected