问题类型

错误

来源

二进制的

tensorflow 版本

文件格式2.10.0

自定义代码

是的

操作系统平台和分发

Windows 10操作系统

移动的设备

• 没有回应 *

Python版本

3.9

Bazel版本

• 没有回应 *

GCC/编译器版本

• 没有回应 *

CUDA/cuDNN版本

工具箱=11.2立方英尺=8.1.0

GPU型号和内存

GTX 1080钛合金

当前行为？

jit_compile=True时此模型的XLA编译器错误
很难确定错误的原因，因为模型在正常运行时工作正常。

重现问题的独立代码

import tensorflow as tf
import numpy as np

class Chooser(tf.keras.layers.Layer):
    @tf.function
    def call(self, options_input, choices_input_logits):
        choices = tf.nn.softmax(choices_input_logits, axis=2)

        result = tf.linalg.matmul(choices, options_input)
        return result

def get_model():
    options_input = tf.keras.layers.Input(shape=(10,3), name="options")
    choices_input = tf.keras.layers.Input(shape=(5,10), name="choices")

    net = Chooser()(options_input, choices_input)
    net = tf.keras.layers.Flatten()(net)
    net = tf.keras.layers.Dense(1)(net)

    return tf.keras.Model([options_input, choices_input], net)

model = get_model()
model.compile(
    optimizer=tf.keras.optimizers.Adam(0.001),
    loss=tf.keras.losses.MeanAbsoluteError(),
    jit_compile=True
)

def batch_gen():
    while True:
        o = np.random.uniform(low=-1.0, high=1.0, size=(10, 3))
        c = np.random.uniform(low=0.0, high=1.0, size=(5, 10))

        y = 1

        yield {"options": o, "choices": c}, y

dataset = tf.data.Dataset.from_generator(batch_gen, output_types=({"options": tf.float32, "choices": tf.float32}, tf.float32))
dataset = dataset.batch(32)

model.fit(dataset, steps_per_epoch=100, epochs=1)

相关日志输出

tensorflow/compiler/xla/client/lib/dynamic_shaped_ops.cc:92] Check failed: left_branch_shape.rank() == right_branch_shape.rank() (1 vs. 2)left rank of (1) vs. right rank of (2)