linear_api_model.py

class Linear(tf.keras.Model):
  def __init__(self):
    super().__init__()
    self.layer = LinearLayer()
  def call(self, input):
    output = self.layer(input)
    return output
  
model = Linear()
optimizer = tf.train.GradientDescentOptimizer(1e-3)

for _ in range(10000):
    with tf.GradientTape() as tape:
        y_pred = model(X)
        loss = tf.reduce_mean(tf.square(y_pred - y))

    grads = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))

print(model.variables)