Add CLIP-JAX ADE20K implementation, training, and tests by Thakor-Yashpal · Pull Request #83 · jax-ml/bonsai

Thakor-Yashpal · 2025-11-12T08:00:50Z

Resolves #<issue_number_goes_here>

Reference

Checklist

[] I have read the Contribution Guidelines and used pre-commit hooks to format this commit.
[] I have added all the necessary unit tests for my change. (run_model.py for model usage, test_outputs.py and/or model_validation_colab.ipynb for quality).
[] (If using an LLM) I have carefully reviewed and removed all superfluous comments or unneeded, commented-out code. Only necessary and functional code remains.
[] I have signed the Contributor License Agreement (CLA).

…ing the unnecessary files.

jenriver · 2025-11-20T19:38:01Z

Hello Thakor, it'd be really cool to have a CLIP implementation here. It seems there's no file changes here, did you forget to push commit?

Related issue: #72

Thakor-Yashpal · 2025-11-21T10:37:19Z

Hello Thakor, it'd be really cool to have a CLIP implementation here. It seems there's no file changes here, did you forget to push commit?

Relaetd issue: #72

I had pushed the initial code earlier, but the implementation wasn’t up to the quality I wanted, so I removed those files. I’m making the necessary improvements now and will push the updated files on Wednesday.

…onsai into clip-jax-ade20k

Iamleos · 2025-12-04T07:56:02Z

When CLIP will be merged?

jenriver · 2025-12-05T23:53:21Z

bonsai/models/clip/run_model.py

@@ -0,0 +1,185 @@
+from typing import Any


Can this be entirely used in tests/ insetad to follow pattern in https://github.com/jax-ml/bonsai/blob/main/CONTRIBUTING.md#contributing-a-model?

chapman20j · 2025-12-11T23:00:04Z

bonsai/models/efficientnet/tests/EfficientNet_ImageNet_validation_example.ipynb

Could you please update this commit to not remove the efficientnet files?

chapman20j · 2025-12-11T23:01:48Z

bonsai/models/clip/run_model.py

+
+def _get_dtype(cfg: CLIPConfig):
+    return jnp.float32 if cfg.dtype == "float32" else jnp.float16
+


It looks like these layers are replicated here. Can we just have the definitions in modeling.py?

chapman20j · 2025-12-11T23:02:16Z

bonsai/models/clip/modeling.py

+import jax.numpy as jnp
+import flax.linen as nn
+from flax.linen import initializers
+from .params import CLIPConfig


Can we have the config in this file for consistency with the rest of the repo?

chapman20j · 2025-12-11T23:03:04Z

bonsai/models/clip/params.py

+            self.text_embed_dim = 1024
+            self.proj_dim = 1024
+        else:
+            raise ValueError("Unknown model_size: " + str(self.model_size))


For users interested in inference, could you add functionality to transfer parameters from a pretrained model?

chapman20j · 2025-12-11T23:06:46Z

Hi @Thakor-Yashpal. Just following up on this PR. I left a few comments but thought I'd ask how things are going. Are you still actively developing CLIP?

Thakor-Yashpal · 2025-12-12T08:01:02Z

Hi @Thakor-Yashpal. Just following up on this PR. I left a few comments but thought I'd ask how things are going. Are you still actively developing CLIP?

Yes, I am! My exams finish tomorrow, and once they're done I’ll be back to working on this PR. I’ll start addressing the comments and updating everything by tomorrow. Thanks for checking in!

chapman20j · 2025-12-16T01:05:19Z

Sounds good. Goodluck with exams!

chapman20j · 2026-01-05T23:06:58Z

Hi @Thakor-Yashpal . I hope your exams went well! Just checking in on this PR. Are you still working on this?

Thakor-Yashpal · 2026-01-06T07:39:36Z

Hi @Thakor-Yashpal . I hope your exams went well! Just checking in on this PR. Are you still working on this?

Hi @chapman20j, thanks for checking in!
Yes, I’m still working on this PR. I’ll make sure everything is updated and submit it by this Friday. No need to remind me again — I’ll definitely have it done by then. Thanks for your patience!

Thakor-Yashpal · 2026-01-12T20:48:59Z

Hi @chapman201, thanks for your patience!
I’ve now submitted the JAX CLIP model implementation along with the forward-pass pytest, and everything is passing on Python 3.11.
Please let me know if you’d like any changes or additional checks. Happy to iterate!

jenriver · 2026-01-12T21:23:13Z

Hi Thakor, could you update the branch and resolve merge conflicts?

jenriver · 2026-01-22T01:18:40Z

Thanks for the contribution! I can't merge this yet due to a few critical blockers:

Restoration Required: You accidentally deleted the efficientnet directory. Please revert those deletions; clip should be a new sibling directory.
Model Correctness: The Transformer layers in modeling.py are missing the MLP/Feed-Forward blocks. Without them, the model is mathematically incorrect.
Missing Features: params.py needs to include logic to load pretrained weights (e.g., from OpenAI/HuggingFace), similar to other models in this repo.
Cleanup: The notebook should import the model classes rather than defining them inline to avoid code drift.

Let me know once these are addressed and I'll take another look!

…onsai into clip-jax-ade20k

Thakor-Yashpal added 2 commits November 12, 2025 13:30

Add CLIP-JAX ADE20K implementation, training, and tests

c4eb01a

I am working on the new model, similar to the previous one, but remov…

5e34692

…ing the unnecessary files.

Thakor-Yashpal added 8 commits November 26, 2025 03:46

Add initial CLIP (ITACLIP) model implementation

d8dbeba

Create README.md

4fded20

Update README.md

f437948

Update README.md

a7a325b

x

98556c2

Delete bonsai/models/clip/tests/__init__.py

2cf8d6e

changes

97cc48d

Merge branch 'clip-jax-ade20k' of https://github.com/Thakor-Yashpal/b…

2d280e9

…onsai into clip-jax-ade20k

jenriver reviewed Dec 5, 2025

View reviewed changes

chapman20j reviewed Dec 11, 2025

View reviewed changes

Thakor-Yashpal added 3 commits January 13, 2026 02:13

Add JAX CLIP model

2350207

.

4c64bf3

.

7c3c29e

Thakor-Yashpal and others added 2 commits January 19, 2026 12:15

Merge remote-tracking branch 'upstream/main' into clip-jax-ade20k

46e0169

Merge branch 'main' into clip-jax-ade20k

d4e9b94

Merge branch 'main' into clip-jax-ade20k

c78dc91

Thakor-Yashpal and others added 3 commits January 22, 2026 13:01

Restore efficientnet directory removed during merge

a9d4018

Merge branch 'clip-jax-ade20k' of https://github.com/Thakor-Yashpal/b…

b2b28a5

…onsai into clip-jax-ade20k

Merge branch 'main' into clip-jax-ade20k

43864df

jenriver marked this pull request as draft January 30, 2026 21:50


		def _get_dtype(cfg: CLIPConfig):
		return jnp.float32 if cfg.dtype == "float32" else jnp.float16

Conversation

Thakor-Yashpal commented Nov 12, 2025

Uh oh!

jenriver commented Nov 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Thakor-Yashpal commented Nov 21, 2025

Uh oh!

Iamleos commented Dec 4, 2025

Uh oh!

jenriver Dec 5, 2025

Choose a reason for hiding this comment

Uh oh!

chapman20j Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

chapman20j Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

chapman20j Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

chapman20j Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

chapman20j commented Dec 11, 2025

Uh oh!

Thakor-Yashpal commented Dec 12, 2025

Uh oh!

chapman20j commented Dec 16, 2025

Uh oh!

chapman20j commented Jan 5, 2026

Uh oh!

Thakor-Yashpal commented Jan 6, 2026

Uh oh!

Thakor-Yashpal commented Jan 12, 2026

Uh oh!

jenriver commented Jan 12, 2026

Uh oh!

jenriver commented Jan 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

jenriver commented Nov 20, 2025 •

edited

Loading