Как узнать, какие тензоры d-типа принимаются данной встроенной функцией тензора потока?

Вы можете это выяснить, но результат будет представлен в терминах ops и ядер, которые не отображаются точно на функции Python более высокого уровня. В случае, если вы не знакомы с архитектурой TensorFlow, она построена вокруг концепции «операций», которые являются просто формальным описанием операции с тензорами (например, операция «Добавить» принимает два значения и выводит третье значение). Граф вычисления TensorFlow состоит из взаимосвязанных операционных узлов. Операции сами по себе не реализуют никакой логики, они просто указывают имя и атрибуты операции, включая типы данных, к которым она может быть применена. Реализация ops дается ядрами, которые являются фактическими частями кода, которые выполняют работу. У одной операции может быть много зарегистрированных ядер, которые работают с разными типами данных и / или разными устройствами (CPU, GPU).

TensorFlow хранит «реестры» со всей этой информацией, хранимой в виде различных сообщений Protocol Buffers . Хотя он не является частью общедоступного API, вы можете запросить эти реестры, чтобы получить список операций или ядер, которые соответствуют определенным критериям. Например, вот как вы можете получить все операции, которые работают с некоторым сложным типом:

import tensorflow as tf

def get_ops_with_dtypes(dtypes):
    from tensorflow.python.framework import ops
    valid_ops = []
    dtype_enums = set(dtype.as_datatype_enum for dtype in dtypes)
    reg_ops = ops.op_def_registry.get_registered_ops()
    for op in reg_ops.values():
        for attr in op.attr:
            if (attr.type == 'type' and
                any(t in dtype_enums for t in attr.allowed_values.list.type)):
                valid_ops.append(op)
                break
    # Sort by name for convenience
    return sorted(valid_ops, key=lambda op: op.name)

complex_dtypes = [tf.complex64, tf.complex128]
complex_ops = get_ops_with_dtypes(complex_dtypes)

# Print one op
print(complex_ops[0])
# name: "AccumulateNV2"
# input_arg {
#   name: "inputs"
#   type_attr: "T"
#   number_attr: "N"
# }
# output_arg {
#   name: "sum"
#   type_attr: "T"
# }
# attr {
#   name: "N"
#   type: "int"
#   has_minimum: true
#   minimum: 1
# }
# attr {
#   name: "T"
#   type: "type"
#   allowed_values {
#     list {
#       type: DT_FLOAT
#       type: DT_DOUBLE
#       type: DT_INT32
#       type: DT_UINT8
#       type: DT_INT16
#       type: DT_INT8
#       type: DT_COMPLEX64
#       type: DT_INT64
#       type: DT_QINT8
#       type: DT_QUINT8
#       type: DT_QINT32
#       type: DT_BFLOAT16
#       type: DT_UINT16
#       type: DT_COMPLEX128
#       type: DT_HALF
#       type: DT_UINT32
#       type: DT_UINT64
#     }
#   }
# }
# attr {
#   name: "shape"
#   type: "shape"
# }
# is_aggregate: true
# is_commutative: true

# Print op names
print(*(op.name for op in complex_ops), sep='\n')
# AccumulateNV2
# AccumulatorApplyGradient
# AccumulatorTakeGradient
# Acos
# Acosh
# Add
# AddN
# AddV2
# Angle
# ApplyAdaMax
# ...

Здесь элементы в complex_ops являются OpDef сообщениями, которые вы можете проверить чтобы узнать точную структуру оп. В этом случае get_ops_with_dtypes просто возвращает каждую операцию, которая имеет один из заданных типов данных среди своих атрибутов type, поэтому комплексное значение может применяться к одному из входов или выходов.

Другой альтернативой может быть непосредственный поиск ядер, работающих с интересующими вас типами данных. Ядра хранятся в виде сообщений KernelDef , которые не содержат всей информации об операциях, но например, у них есть информация об устройстве, на котором они могут работать, поэтому вы также можете запросить ядра, поддерживающие определенное устройство.

import tensorflow as tf

def get_kernels_with_dtypes(dtypes, device_type=None):
    from tensorflow.python.framework import kernels
    valid_kernels = []
    dtype_enums = set(dtype.as_datatype_enum for dtype in dtypes)
    reg_kernels = kernels.get_all_registered_kernels()
    for kernel in reg_kernels.kernel:
        if device_type and kernel.device_type != device_type:
            continue
        for const in kernel.constraint:
            if any(t in dtype_enums for t in const.allowed_values.list.type):
                valid_kernels.append(kernel)
                break
    # Sort by name for convenience
    return sorted(valid_kernels, key=lambda kernel: kernel.op)

complex_dtypes = [tf.complex64, tf.complex128]
complex_gpu_kernels = get_kernels_with_dtypes(complex_dtypes, device_type='GPU')

# Print one kernel
print(complex_gpu_kernels[0])
# op: "Add"
# device_type: "GPU"
# constraint {
#   name: "T"
#   allowed_values {
#     list {
#       type: DT_COMPLEX64
#     }
#   }
# }

# Print kernel op names
print(*(kernel.op for kernel in complex_gpu_kernels), sep='\n')
# Add
# Add
# AddN
# AddN
# AddV2
# AddV2
# Assign
# Assign
# AssignVariableOp
# AssignVariableOp
# ...

Проблема в том, что вы никогда не используете ops или ядра напрямую, когда программируете с TensorFlow на Python. Функции Python принимают аргументы, которые вы им даете, проверяете их и производите один или несколько новых операций на графике, обычно возвращая вам выходные значения последнего из них. Поэтому, в конце концов, выяснение того, какие ops / ядра вам подходят, требует небольшой проверки. Например, рассмотрим следующие примеры:

import tensorflow as tf

with tf.Graph().as_default():
    # Matrix multiplication: (2, 3) x (3, 4)
    tf.matmul([[1, 2, 3], [4, 5, 6]], [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
    # Print all op names and types
    all_ops = tf.get_default_graph().get_operations()
    print(*(f'Op name: {op.name}, Op type: {op.type}' for op in all_ops), sep='\n')
    # Op name: MatMul/a, Op type: Const
    # Op name: MatMul/b, Op type: Const
    # Op name: MatMul, Op type: MatMul

with tf.Graph().as_default():
    # Matrix multiplication: (1, 2, 3) x (1, 3, 4)
    tf.matmul([[[1, 2, 3], [4, 5, 6]]], [[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]])
    # Print all op names and types
    all_ops = tf.get_default_graph().get_operations()
    print(*(f'Op name: {op.name}, Op type: {op.type}' for op in all_ops), sep='\n')
    # Op name: MatMul/a, Op type: Const
    # Op name: MatMul/b, Op type: Const
    # Op name: MatMul, Op type: BatchMatMul

Здесь та же самая функция Python tf.matmul выдает типы операций в каждом случае. Первые два операции - Const в обоих случаях, что является результатом преобразования данных списков в тензоры TensorFlow, но третий - MatMul в одном случае и BatchedMatMul в другом, потому что во втором случае вход имеет одно дополнительное начальное измерение.

В любом случае, если вы можете объединить описанные выше подходы, чтобы узнать всю информацию о ядрах и операциях об одном интересующем вас имени операции:

def get_op_info(op_name):
    from tensorflow.python.framework import ops
    from tensorflow.python.framework import kernels
    reg_ops = ops.op_def_registry.get_registered_ops()
    op_def = reg_ops[op_name]
    op_kernels = list(kernels.get_registered_kernels_for_op(op_name).kernel)
    return op_def, op_kernels

# Get MatMul information
matmul_def, matmul_kernels = get_op_info('MatMul')

# Print op definition
print(matmul_def)
# name: "MatMul"
# input_arg {
#   name: "a"
#   type_attr: "T"
# }
# input_arg {
#   name: "b"
#   type_attr: "T"
# }
# output_arg {
#   name: "product"
#   type_attr: "T"
# }
# attr {
#   name: "transpose_a"
#   type: "bool"
#   default_value {
#     b: false
#   }
# }
# attr {
#   name: "transpose_b"
#   type: "bool"
#   default_value {
#     b: false
#   }
# }
# attr {
#   name: "T"
#   type: "type"
#   allowed_values {
#     list {
#       type: DT_BFLOAT16
#       type: DT_HALF
#       type: DT_FLOAT
#       type: DT_DOUBLE
#       type: DT_INT32
#       type: DT_COMPLEX64
#       type: DT_COMPLEX128
#     }
#   }
# }

# Total number of matrix multiplication kernels
print(len(matmul_kernels))
# 24

# Print one kernel definition
print(matmul_kernels[0])
# op: "MatMul"
# device_type: "CPU"
# constraint {
#   name: "T"
#   allowed_values {
#     list {
#       type: DT_FLOAT
#     }
#   }
# }