Python Keras 不会将模型输出与为整个小批量设计的掩码相乘

我有一个数据生成器，可以生成成批的输入数据(X)和目标(Y)，还有一个应用于模型输出的掩码(batch_mask)(同一个掩码适用于该批中的所有数据点；不同的批次有不同的掩码，数据生成器负责执行此操作).

因此，batch_mask的第一维度可以具有形状1或batch_size(通过沿第一维度重复相同遮罩batch_size次).我希望Keras允许我使用其中一种，我只想在第一维度上创建形状为1的面具.

然而，当我try 这个时，我得到了一个错误:

ValueError: Data cardinality is ambiguous:
  x sizes: 128, 1
  y sizes: 128
Make sure all arrays contain the same number of samples.

这似乎并不复杂.

下面是一些观察这种行为的简单示例代码

import tensorflow.keras as tfk
import numpy as np

#######################
# 1. model definition #
#######################

# model parameters
nfeatures_in = 6
target_size = 8

# model inputs
input = tfk.layers.Input(nfeatures_in)
input_mask = tfk.layers.Input(target_size)

# model graph
out = tfk.layers.Dense(target_size)(input)
out_masked = tfk.layers.Multiply()((out,input_mask)) # multiply all model outputs in the batch by the same mask
model = tfk.Model(inputs=(input, input_mask), outputs=out_masked)

##########################
# 2. dummy data creation #
##########################

batch_size = 32

# create masks the batch
zeros_vector = np.zeros((1,target_size)) # "batch_size"==1
zeros_vector[0,:6] = 1
batch_mask = zeros_vector

# dummy data creation
X = np.random.randn(batch_size, 6)
Y = np.random.randn(batch_size, target_size)*batch_mask # the target is masked by design in each batch


############################
# 3. compile model and fit #
############################

model.compile(optimizer="Adam", loss="mse")
model.fit((X, batch_mask),Y, batch_size=batch_size)

我知道我可以通过以下两种方式来实现这一目标:

• 重复遮罩，使第一个维度batch_mask的大小为第一个维度X的大小(而不是1).
• 使用纯tensorflow(但我觉得对Keras来说，沿着批处理维度进行广播应该不是问题).

How can I make this work with Keras?

非常感谢.