Добавление новой Op в Tensorflow: функции формы

Question

Я пытаюсь добавить новую Op (используя этот способ: https://www.tensorflow.org/how_tos/adding_an_op/). У меня нет проблем с простой Op, как в примере, но мне нужно добавить более сложную Op. Он должен иметь 2 входа и выход, имеющий ту же форму, что и матричное умножение входов. Как написать функцию формы для этого случая? Как я могу выделить вывод с соответствующей формой? Спасибо заранее.

Answer 1

Имейте частичный ответ на вопрос. я до сих пор не знаю, как написать функцию формы, но я выписывал чеки, как это:

const Tensor& input1 = context->input(0); 
const Tensor& input2 = context->input(1); 
TensorShape sh1 = input1.shape(); 
TensorShape sh2 = input2.shape(); 
OP_REQUIRES(context, sh1.dim_size(1)==sh2.dim_size(0), 
errors::InvalidArgument("Can't multiplicate!")); 

Выделяя:

sh1.RemoveDim(1); 
sh2.RemoveDim(0); 
sh1.AppendShape(sh2); 
OP_REQUIRES_OK(context, context->allocate_output(0, sh1, &output_tensor)); 

Но, кажется, что я заново изобретать велосипед. Легче ли это сделать?