Как разбить список входных данных в тенорпотоках керас Модель вызова метода?

Question

Как мы должны разделить входы? Получение различных ошибок. tf.split похоже, что он должен работать, но не уверен, что это правильно?

import numpy as np
import pandas as pd
import tensorflow as tf
import tensorflow.keras as keras

tf.keras.backend.set_floatx('float64')

m = 250  # samples
n_x = 1  # dim of x
n_tau = 11

x = (2 * np.random.rand(m, n_x).astype(np.float64) - 1) * 2 
i = np.argsort(x[:, 0]) 
x = x[i]  # to make plotting nicer
A = np.random.randn(n_x, 1)
y = x ** 2 + 0.3 * x + 0.4 * np.random.randn(m, 1).astype(np.float64)
y = y.dot(A)  # y is 1d
# y = y.squeeze()
tau = np.linspace(1.0 / n_tau, 1 - 1.0 / n_tau, n_tau).astype(np.float64)
tau = tau[:, None]

def loss(tau_y, u):
    tau, y = tau_y
    tf.debugging.assert_rank(y, 2, f"y should be rank 2")
    u = y[:, None, :] - u[None, :, :]
    tf.debugging.assert_rank(tau, 2, f"tau should be rank 2")
    tau = tau[None, :, :]
    res = u ** 2 * (tau - tf.where(u <= np.float64(0.0), np.float64(1.0), np.float64(0.0)))
    return tf.reduce_sum(tf.reduce_mean(res, axis=[1, 2]), axis=0)


tf.keras.backend.set_floatx('float64')
class My(tf.keras.models.Model):
   def __init__(self):
       super().__init__()
       self._my_layer = tf.keras.layers.Dense(1, dtype=tf.float64)
   def call(self, inputs):
       tau, y = inputs
       tf.print(tau.shape, y.shape)
       return self._my_layer(tau)


model = My()
u = model((tau, y))
loss((tau, y), model((tau, y)))
opt = tf.keras.optimizers.Adam(learning_rate=0.01)
model.compile(optimizer=opt, loss=loss)
model.fit((tau, y), (tau, y)) # fails

Это работает для вызова, но не из скомпилированной модели. Какая-то странная ошибка в понимании списка.

~/j.py in loss(tau_y, u)
     21 
     22 def loss(tau_y, u):
---> 23     tau, y = tau_y
     24     tf.debugging.assert_rank(y, 2, f"y should be rank 2")
     25     u = y[:, None, :] - u[None, :, :]

~/anaconda3/envs/37nightly/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py in __iter__(self)
    537   def __iter__(self):
    538     if not context.executing_eagerly():
--> 539       self._disallow_iteration()
    540 
    541     shape = self._shape_tuple()

~/anaconda3/envs/37nightly/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py in _disallow_iteration(self)
    533     else:
    534       # Default: V1-style Graph execution.
--> 535       self._disallow_in_graph_mode("iterating over `tf.Tensor`")
    536 
    537   def __iter__(self):

~/anaconda3/envs/37nightly/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py in _disallow_in_graph_mode(self, task)
    513     raise errors.OperatorNotAllowedInGraphError(
    514         "{} is not allowed in Graph execution. Use Eager execution or decorate"
--> 515         " this function with @tf.function.".format(task))
    516 
    517   def _disallow_bool_casting(self):

OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed in Graph execution. Use Eager execution or decorate this function with @tf.function.