Fix tensorflow transformer

[dansola]

Tracking issue

flyteorg/flyte#6427

Why are the changes needed?

The TensorFlow model transformer was missing the .keras file extension, causing Keras 3 to fail with "File format not supported".

What changes were proposed in this pull request?

Fixed by adding .keras extension to save/load paths and switching from MULTIPART to SINGLE blob dimensionality to match other ML transformers (PyTorch, Sklearn).

How was this patch tested?

Tested on a Union cluster here: https://demo.hosted.unionai.cloud/console/projects/flytesnacks/domains/development/executions/a6vjhxr5hsd75tmdqsfc/nodes

import flytekit as fl
import tensorflow as tf

new_flytekit = "git+https://github.com/flyteorg/flytekit@661edad5ea6b91d2abd4129535a80e938db01cb4"
image = fl.ImageSpec(
    builder="union",
    apt_packages=["git"],
    packages=["union", "tensorflow", "tensorflow-datasets", "flytekitplugins-kftensorflow", new_flytekit],
)

@fl.task(container_image=image)
def train_model() -> tf.keras.Model:
    model = tf.keras.Sequential(
        [tf.keras.layers.Dense(128, activation="relu"), tf.keras.layers.Dense(10, activation="softmax")]
    )
    model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
    return model


@fl.task(container_image=image)
def evaluate_model(model: tf.keras.Model, x: tf.Tensor, y: tf.Tensor) -> float:
    loss, accuracy = model.evaluate(x, y)
    return accuracy


@fl.workflow
def training_workflow() -> float:
    # Hardcoded tensors for testing
    x = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
    y = tf.constant([0, 1, 2, 3])

    model = train_model()
    return evaluate_model(model=model, x=x, y=y)

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

Summary by Bito

• Changes the blob dimensionality from MULTIPART to SINGLE in the TensorFlow model transformer to ensure compatibility with Keras 3.
• Adds the .keras file extension to save/load paths for the TensorFlow model transformer.
• Updates relevant unit tests to validate the new behavior of the TensorFlow model transformer.
• Overall summary: modifies TensorFlow model transformer configurations and updates unit tests to ensure compatibility with Keras 3.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 46.50%. Comparing base (9f3ec96) to head (661edad).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3346      +/-   ##
==========================================
+ Coverage   44.28%   46.50%   +2.21%     
==========================================
  Files         305      216      -89     
  Lines       27188    22615    -4573     
  Branches     2970     2970              
==========================================
- Hits        12040    10516    -1524     
+ Misses      15051    11577    -3474     
- Partials       97      522     +425     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.