Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/framework/byte_swap_tensor.py: 24%

1# Copyright 2023 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================== 

15 

16"""Utilities for byte swapping the tensor content.""" 

17 

18from tensorflow.core.framework import graph_pb2 

19from tensorflow.core.protobuf import meta_graph_pb2 

20from tensorflow.python.framework import dtypes 

21 

22# Based on tensor_bundle/byte_swap.cc 

23byte_swappable = [ 

24 dtypes.float16, 

25 dtypes.float32, 

26 dtypes.float64, 

27 dtypes.bfloat16, 

28 dtypes.complex64, 

29 dtypes.complex128, 

30 dtypes.uint16, 

31 dtypes.uint32, 

32 dtypes.uint64, 

33 dtypes.int16, 

34 dtypes.int32, 

35 dtypes.int64, 

36 dtypes.qint16, 

37 dtypes.quint16, 

38 dtypes.qint32, 

39] 

40 

41 

42def byte_swap_tensor_content(tensor, from_endiness, to_endiness): 

43 """Byte swaps. 

44 

45 Args: 

46 tensor: Target tensor to change endiness. 

47 from_endiness: The original endianness format. "big" or "little" 

48 to_endiness: The target endianness format. "big" or "little" 

49 """ 

50 if tensor.dtype in byte_swappable: 

51 tshape = tensor.tensor_shape.dim 

52 tensor_bytes = tensor.tensor_content 

53 if tensor_bytes: 

54 tensor_size = 1 

55 for sz in tshape: 

56 if sz.size != 0: 

57 tensor_size = tensor_size * sz.size 

58 chunksize = int(len(tensor_bytes) / tensor_size) 

59 # Split tensor_data into chunks for byte swapping. 

60 to_swap = [ 

61 tensor_bytes[i : i + chunksize] 

62 for i in range(0, len(tensor_bytes), chunksize) 

63 ] 

64 # Swap and replace tensor_content. 

65 tensor.tensor_content = b"".join( 

66 [ 

67 int.from_bytes(byteswap, from_endiness).to_bytes( 

68 chunksize, to_endiness 

69 ) 

70 for byteswap in to_swap 

71 ] 

72 ) 

73 

74 

75def swap_tensor_content_in_graph_function( 

76 graph_def, from_endiness, to_endiness 

77): 

78 """Fix endiness of tensor contents. 

79 

80 Args: 

81 graph_def: Target graph_def to change endiness. 

82 from_endiness: The original endianness format. "big" or "little" 

83 to_endiness: The target endianness format. "big" or "little" 

84 """ 

85 if isinstance(graph_def, meta_graph_pb2.MetaGraphDef): 

86 functions = graph_def.graph_def.library.function 

87 elif isinstance(graph_def, graph_pb2.GraphDef): 

88 functions = graph_def.library.function 

89 else: 

90 return 

91 for function in functions: 

92 node_def = function.node_def 

93 for node in node_def: 

94 if node.op == "Const": 

95 tensor = node.attr["value"].tensor 

96 byte_swap_tensor_content(tensor, from_endiness, to_endiness) 

97 

98 

99def swap_tensor_content_in_graph_node(graph_def, from_endiness, to_endiness): 

100 for node in graph_def.node: 

101 if node.op == "Const": 

102 tensor = node.attr["value"].tensor 

103 byte_swap_tensor_content(tensor, from_endiness, to_endiness)