text: # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the 'License');
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an 'AS IS' BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ======================================
  """Utilities for XLA-specific Python types."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import collections
  
  import numpy as _np  # Avoids becoming a part of public Tensorflow API.
  
  from tensorflow.compiler.xla import xla_data_pb2
  from tensorflow.python.framework import dtypes
  
  # Records corresponsence between a XLA primitive type and Python/Numpy types.
  #
  # primitive_type: value of type xla_data_pb2.PrimitiveType
  # numpy_dtype: corresponsing Numpy "dtype" (like np.float32)
  # literal_field_name: name of the field in the LiteralProto message elements
  # of this type go into.
  # literal_field_type: type of the field named 'literal_field_name'.
  #
  # TODO(eliben): figure out how to avoid knowing the extra Python type and the
  # astype cast when writing into Literals.
  TypeConversionRecord = collections.namedtuple('TypeConversionRecord', [
      'primitive_type', 'numpy_dtype', 'literal_field_name', 'literal_field_type'
  ])
  
  # Maps from XLA primitive types to TypeConversionRecord.
  MAP_XLA_TYPE_TO_RECORD = {
      xla_data_pb2.BF16:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.BF16,
              numpy_dtype=dtypes.bfloat16.as_numpy_dtype,
              literal_field_name='bf16s',
              literal_field_type=float),
      xla_data_pb2.F16:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.F16,
              numpy_dtype=_np.float16,
              literal_field_name='f16s',
              literal_field_type=float),
      xla_data_pb2.F32:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.F32,
              numpy_dtype=_np.float32,
              literal_field_name='f32s',
              literal_field_type=float),
      xla_data_pb2.F64:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.F64,
              numpy_dtype=_np.float64,
              literal_field_name='f64s',
              literal_field_type=float),
      xla_data_pb2.S8:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.S8,
              numpy_dtype=_np.int8,
              literal_field_name='s8s',
              literal_field_type=int),
      xla_data_pb2.S16:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.S16,
              numpy_dtype=_np.int16,
              literal_field_name='s16s',
              literal_field_type=int),
      xla_data_pb2.S32:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.S32,
              numpy_dtype=_np.int32,
              literal_field_name='s32s',
              literal_field_type=int),
      xla_data_pb2.S64:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.S64,
              numpy_dtype=_np.int64,
              literal_field_name='s64s',
              literal_field_type=int),
      xla_data_pb2.U8:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.U8,
              numpy_dtype=_np.uint8,
              literal_field_name='s8s',
              literal_field_type=int),
      xla_data_pb2.U16:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.U16,
              numpy_dtype=_np.uint16,
              literal_field_name='s16s',
              literal_field_type=int),
      xla_data_pb2.U32:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.U32,
              numpy_dtype=_np.uint32,
              literal_field_name='s32s',
              literal_field_type=int),
      xla_data_pb2.U64:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.U64,
              numpy_dtype=_np.uint64,
              literal_field_name='s64s',
              literal_field_type=int),
      xla_data_pb2.PRED:
          TypeConversionRecord(
              primitive_type=xla_data_pb2.PRED,
              numpy_dtype=_np.bool,
              literal_field_name='preds',
              literal_field_type=bool)
  }
  
  # Maps from Numpy dtypes to TypeConversionRecord.
  # Note the conversion on the key. Numpy has a known issue wherein dtype hashing
  # doesn't work as expected (https://github.com/numpy/numpy/issues/7242). Thus,
  # when keying by dtype in this dict, we use the string form of dtypes.
  MAP_DTYPE_TO_RECORD = {
      str(_np.dtype(record.numpy_dtype)): record
      for record in MAP_XLA_TYPE_TO_RECORD.values()
  }
  