DistributionLambda incompatible with graph mode

From this documentation, TFP can be used in graph mode. However, in the following code where I trained a simple probabilistic NN regression (by using DistributionLambda layer) in graph mode, the prediction of the model is not interpretable.

import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf
from tensorflow.python.eager.context import eager_mode, graph_mode

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras import Input
from tensorflow.keras.initializers import RandomNormal
from tensorflow.keras.optimizers import Adam

import tensorflow_probability as tfp

tfd = tfp.distributions

from datetime import datetime

# Do not use GPU
tf.config.set_visible_devices([], "GPU")

# Generate data
def f(x, noise_level):
    Generator = np.random.default_rng()
    noise = Generator.normal(0, noise_level, x.shape)
    return (x + 1) * np.sin(5 * x) + noise


x_plot = np.arange(-1, 1 + 0.001, 0.001)
y_plot = f(x_plot, 0)

x_train = np.arange(-1 + 0.05, 1, 0.2)
y_train = f(x_train, 0.05)

x_val = np.arange(-1 + 0.15, 1, 0.2)
y_val = f(x_val, 0.05)

# Plot the problem
plt.figure()
plt.plot(x_plot, y_plot, "-", label="Orgininal function without noise")
plt.plot(x_train, y_train, "o", label="Training points")
plt.plot(x_val, y_val, "s", label="Validation points")
plt.xlim(-1, 1)
plt.ylim(-2, 2)
plt.xlabel("x")
plt.ylabel("f")
plt.grid()
plt.legend()
plt.show(block=False)

# Reshape
X_train = x_train.reshape(x_train.shape[0], 1)
Y_train = y_train.reshape(x_train.shape[0], 1)

X_val = x_val.reshape(x_val.shape[0], 1)
Y_val = y_val.reshape(x_val.shape[0], 1)

# Test model
start_time = datetime.now()


def train():
    tf.keras.utils.set_random_seed(1)

    def negloglik(y, rv_y):
        return -rv_y.log_prob(y)

    model = Sequential()
    model.add(Input(shape=(1,)))  # Input layer
    model.add(
        Dense(
            4,
            activation="sigmoid",
            kernel_initializer=RandomNormal(mean=0.0, stddev=1.0),
        )
    )
    model.add(
        Dense(
            2,
            kernel_initializer=RandomNormal(mean=0.0, stddev=1.0),
        )
    )
    model.add(
        tfp.layers.DistributionLambda(
            lambda t: tfd.Normal(
                loc=t[..., :1],
                scale=1e-3 + tf.math.softplus(0.05 * t[..., 1:]),
            )
        )
    )
    model.compile(
        loss=negloglik,
        optimizer=Adam(learning_rate=3e-2),
        run_eagerly=False,
    )
    history = model.fit(
        X_train,
        Y_train,
        validation_split=0.0,
        validation_data=(X_val, Y_val),
        validation_freq=1,
        batch_size=X_train.shape[0],
        epochs=1000,
        verbose=0,
    )
    return model


model = train()

run_time = datetime.now() - start_time
print("Training time : {:.4f} s".format(run_time.total_seconds()))

# with eager_mode():
with graph_mode():
    model = train()
    pred = model(X_val)
    mean = pred.mean()

print(mean)

mean = mean.numpy()

The printed tensor output is strange to me:

Tensor("tensor_coercible_CONSTRUCTED_AT_sequential_15_distribution_lambda_15/mean/mul:0", shape=(10, 1), dtype=float32)

I cannot converge this tensor to an numpy array for other use, plotting for example. This do not happen when I build and train the model in eager mode. How can I converge this tensor to normal tensor which looks like it:

tf.Tensor(
[[ 0.10509682]
 [ 0.00954151]
 [-0.46975946]
 [-0.76375914]
 [-0.20615435]
 [ 0.74467945]
 [ 1.2763762 ]
 [ 0.61059904]
 [-1.0483716 ]
 [-2.0227783 ]], shape=(10, 1), dtype=float32)

I am using tf==2.9.0 and tfp=0.17.0. Thank you!

Start with the supplied train() reproduction and compare the graph_mode and eager_mode paths around model(X_val), pred.mean(), and mean.numpy(). Investigate how DistributionLambda predictions are represented in graph mode and whether they can be materialized for plotting. Done means graph-mode prediction output is usable as NumPy values, or the limitation and supported alternative are documented.