def bfloat16_to_float32_nested(input_nested):
  """Convert float32 tensors in a nested structure to bfloat16.

  Args:
    input_nested: A Python dict, values being Tensor or Python list/tuple of
      Tensor or Non-Tensor.

  Returns:
    A Python dict with the same structure as `tensor_dict`,
    with all bfloat16 tensors converted to float32.
  """
  if isinstance(input_nested, tf.Tensor):
    if input_nested.dtype == tf.bfloat16:
      return tf.cast(input_nested, dtype=tf.float32)
    else:
      return input_nested
  elif isinstance(input_nested, (list, tuple)):
    out_tensor_dict = [bfloat16_to_float32_nested(t) for t in input_nested]
  elif isinstance(input_nested, dict):
    out_tensor_dict = {
        k: bfloat16_to_float32_nested(v) for k, v in input_nested.items()
    }
  else:
    return input_nested
  return out_tensor_dict
Convert float32 tensors in a nested structure to bfloat16.

  Args:
    input_nested: A Python dict, values being Tensor or Python list/tuple of
      Tensor or Non-Tensor.

  Returns:
    A Python dict with the same structure as `tensor_dict`,
    with all bfloat16 tensors converted to float32.
  
inference_api/src/main/object_detection/utils/ops.py