Migration of Supervised contrative Learning tutorial to Keras3 by maitry63 · Pull Request #2398 · keras-team/keras-io

maitry63 · 2026-07-17T07:02:44Z

This PR is a follow-up to the work done in #2338 that migrated tutorial Supervised contrative Learning into Keras3.

Added the required generated files.
Fixed the related test cases.
Completed the remaining follow-up changes.

gemini-code-assist

Code Review

This pull request migrates the supervised contrastive learning example to Keras 3, replacing TensorFlow and TensorFlow Addons dependencies with unified Keras APIs and keras.ops. The review feedback suggests several improvements to the custom SupervisedContrastiveLoss implementation across the Python script, Jupyter Notebook, and Markdown documentation. Specifically, it recommends implementing the get_config method for proper serialization, using dynamic dtypes (logits.dtype) instead of hardcoded "float32" to ensure mixed-precision compatibility, and replacing the non-idiomatic ops.subtract(0.0, ...) with a simple unary negation operator.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-17T07:04:48Z

+    def __init__(self, temperature=0.05, **kwargs):
+        super().__init__(**kwargs)
        self.temperature = temperature


To support proper serialization and deserialization of the custom loss (e.g., when saving and loading the model), it is best practice in Keras 3 to implement the get_config method.

def __init__(self, temperature=0.05, **kwargs): super().__init__(**kwargs) self.temperature = temperature def get_config(self): config = super().get_config() config.update({"temperature": self.temperature}) return config

gemini-code-assist · 2026-07-17T07:04:48Z

+        mask = ops.cast(ops.equal(labels, ops.transpose(labels)), "float32")
+
+        batch_size = ops.shape(logits)[0]
+        logits_mask = 1.0 - ops.eye(batch_size)


To ensure full compatibility with mixed precision training (e.g., mixed_float16 or mixed_bfloat16) and to prevent potential runtime dtype mismatch errors in strict backends like PyTorch or JAX, it is highly recommended to use the dynamic dtype of the input logits (i.e., logits.dtype) instead of hardcoding "float32" and default float for ops.eye.

Suggested change

mask = ops.cast(ops.equal(labels, ops.transpose(labels)), "float32")

batch_size = ops.shape(logits)[0]

logits_mask = 1.0 - ops.eye(batch_size)

mask = ops.cast(ops.equal(labels, ops.transpose(labels)), logits.dtype)

batch_size = ops.shape(logits)[0]

logits_mask = 1.0 - ops.eye(batch_size, dtype=logits.dtype)

gemini-code-assist · 2026-07-17T07:04:48Z

+            ops.sum(mask, axis=1) + 1e-8
+        )
+
+        return ops.subtract(0.0, ops.mean(mean_log_prob_pos))


Using ops.subtract(0.0, ...) is non-idiomatic and less readable. We can simply use the unary negation operator - on the tensor, which is fully supported across all Keras backends.

Suggested change

return ops.subtract(0.0, ops.mean(mean_log_prob_pos))

return -ops.mean(mean_log_prob_pos)

gemini-code-assist · 2026-07-17T07:04:48Z

+    "    def __init__(self, temperature=0.05, **kwargs):\n",
+    "        super().__init__(**kwargs)\n",
    "        self.temperature = temperature\n",


To support proper serialization and deserialization of the custom loss (e.g., when saving and loading the model), it is best practice in Keras 3 to implement the get_config method.

Suggested change

" def __init__(self, temperature=0.05, **kwargs):\n",

" super().__init__(**kwargs)\n",

" self.temperature = temperature\n",

" def __init__(self, temperature=0.05, **kwargs):\n",

" super().__init__(**kwargs)\n",

" self.temperature = temperature\n",

"\n",

" def get_config(self):\n",

" config = super().get_config()\n",

" config.update({\"temperature\": self.temperature})\n",

" return config\n",

gemini-code-assist · 2026-07-17T07:04:48Z

+    "        mask = ops.cast(ops.equal(labels, ops.transpose(labels)), \"float32\")\n",
+    "\n",
+    "        batch_size = ops.shape(logits)[0]\n",
+    "        logits_mask = 1.0 - ops.eye(batch_size)\n",


To ensure full compatibility with mixed precision training (e.g., mixed_float16 or mixed_bfloat16) and to prevent potential runtime dtype mismatch errors in strict backends like PyTorch or JAX, it is highly recommended to use the dynamic dtype of the input logits (i.e., logits.dtype) instead of hardcoding "float32" and default float for ops.eye.

Suggested change

" mask = ops.cast(ops.equal(labels, ops.transpose(labels)), \"float32\")\n",

"\n",

" batch_size = ops.shape(logits)[0]\n",

" logits_mask = 1.0 - ops.eye(batch_size)\n",

" mask = ops.cast(ops.equal(labels, ops.transpose(labels)), logits.dtype)\n",

"\n",

" batch_size = ops.shape(logits)[0]\n",

" logits_mask = 1.0 - ops.eye(batch_size, dtype=logits.dtype)\n",

gemini-code-assist · 2026-07-17T07:04:48Z

+    "            ops.sum(mask, axis=1) + 1e-8\n",
+    "        )\n",
+    "\n",
+    "        return ops.subtract(0.0, ops.mean(mean_log_prob_pos))\n",


Using ops.subtract(0.0, ...) is non-idiomatic and less readable. We can simply use the unary negation operator - on the tensor, which is fully supported across all Keras backends.

Suggested change

" return ops.subtract(0.0, ops.mean(mean_log_prob_pos))\n",

" return -ops.mean(mean_log_prob_pos)\n"

gemini-code-assist · 2026-07-17T07:04:48Z

+    def __init__(self, temperature=0.05, **kwargs):
+        super().__init__(**kwargs)
        self.temperature = temperature


To support proper serialization and deserialization of the custom loss (e.g., when saving and loading the model), it is best practice in Keras 3 to implement the get_config method.

Suggested change

def __init__(self, temperature=0.05, **kwargs):

super().__init__(**kwargs)

self.temperature = temperature

def __init__(self, temperature=0.05, **kwargs):

super().__init__(**kwargs)

self.temperature = temperature

def get_config(self):

config = super().get_config()

config.update({"temperature": self.temperature})

return config

gemini-code-assist · 2026-07-17T07:04:48Z

+        mask = ops.cast(ops.equal(labels, ops.transpose(labels)), "float32")
+
+        batch_size = ops.shape(logits)[0]
+        logits_mask = 1.0 - ops.eye(batch_size)


To ensure full compatibility with mixed precision training (e.g., mixed_float16 or mixed_bfloat16) and to prevent potential runtime dtype mismatch errors in strict backends like PyTorch or JAX, it is highly recommended to use the dynamic dtype of the input logits (i.e., logits.dtype) instead of hardcoding "float32" and default float for ops.eye.

Suggested change

mask = ops.cast(ops.equal(labels, ops.transpose(labels)), "float32")

batch_size = ops.shape(logits)[0]

logits_mask = 1.0 - ops.eye(batch_size)

mask = ops.cast(ops.equal(labels, ops.transpose(labels)), logits.dtype)

batch_size = ops.shape(logits)[0]

logits_mask = 1.0 - ops.eye(batch_size, dtype=logits.dtype)

gemini-code-assist · 2026-07-17T07:04:48Z

+            ops.sum(mask, axis=1) + 1e-8
+        )
+
+        return ops.subtract(0.0, ops.mean(mean_log_prob_pos))


Using ops.subtract(0.0, ...) is non-idiomatic and less readable. We can simply use the unary negation operator - on the tensor, which is fully supported across all Keras backends.

Suggested change

return ops.subtract(0.0, ops.mean(mean_log_prob_pos))

return -ops.mean(mean_log_prob_pos)

hertschuh · 2026-07-21T17:34:27Z

-from tensorflow.keras import layers
+import os
+
+os.environ["KERAS_BACKEND"] = "tensorflow"  # or "torch" or "jax"


Any reason not to use "jax"?

hertschuh

Can you test with the JAX backend just to make sure?

maitry63 · 2026-08-04T11:59:27Z

Can you test with the JAX backend just to make sure?

Sure, tested with Jax backend.

github-actions · 2026-08-19T01:58:49Z

This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

hertschuh

There is a merge conflict:

hertschuh · 2026-08-28T21:30:02Z

+    "keras": f"{KERAS_TEAM_GH}/keras/tree/v3.15.1/",
    "keras_tuner": f"{KERAS_TEAM_GH}/keras-tuner/tree/v1.4.8/",
    "keras_hub": f"{KERAS_TEAM_GH}/keras-hub/tree/v0.30.0/",
-    "tf_keras": f"{KERAS_TEAM_GH}/tf-keras/tree/v2.20.0/",
+    "tf_keras": f"{KERAS_TEAM_GH}/tf-keras/tree/v2.20.1/",


Can you rebase? I believe these would go away.

Rebased, to fix the merge conflicts.

… model

github-actions Bot assigned laxmareddyp Jul 17, 2026

gemini-code-assist Bot reviewed Jul 17, 2026

View reviewed changes

maitry63 mentioned this pull request Jul 21, 2026

Migration of Supervised contrative Learning tutorial to Keras3 #2338

Closed

hertschuh reviewed Jul 21, 2026

View reviewed changes

hertschuh reviewed Aug 3, 2026

View reviewed changes

hertschuh added the stat:awaiting response from contributor label Aug 3, 2026

maitry63 force-pushed the sup_contrast_br-followup branch from e2d38eb to 3e59b2a Compare August 4, 2026 11:55

github-actions Bot added the stale label Aug 19, 2026

maitry63 removed the stale label Aug 19, 2026

hertschuh requested changes Aug 28, 2026

View reviewed changes

maitry63 force-pushed the sup_contrast_br-followup branch from 641cd91 to 05a24a9 Compare August 31, 2026 05:52

LakshmiKalaKadali and others added 7 commits August 31, 2026 11:44

Migration of Supervised contrative Learning tutorial to Keras3

9d8778a

Run the code on Jax backend

bac31f4

Removed commentline for training statement in baseline classification…

742c241

… model

Add Corresponding files for Supervised Contrastive Learning Tutorial

94c2e0b

add tested files with jax backend

e21ca5b

add tested files with jax backend

748cbfd

fix: testcases

5d79f58

maitry63 force-pushed the sup_contrast_br-followup branch from 05a24a9 to 5d79f58 Compare August 31, 2026 06:17

	return ops.subtract(0.0, ops.mean(mean_log_prob_pos))
	return -ops.mean(mean_log_prob_pos)

	" return ops.subtract(0.0, ops.mean(mean_log_prob_pos))\n",
	" return -ops.mean(mean_log_prob_pos)\n"

Uh oh!

Conversation

maitry63 commented Jul 17, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

hertschuh Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

maitry63 commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 19, 2026

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

hertschuh Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

maitry63 Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants