new: Implement ModernBERT (#7)

* new: RoPE implementation and tests

* make: apply linting with ruff

* docs: add link to reference implementation for RoPE

* make: fix type checking with mypy

* new: add embedding and MLP layers for modernbert

* new: add ModernBERTAttention

* tests: add numerical tests for ModernBERTAttention, apply linting with Ruff

* make: linting and proper type annotation

* new: add ModerBERTEncoderLayer

* make: fix type hints

* make: better type annotations

* new: add full ModernBERT model implementation

* docs: update docstrings for ModernBERT

* make: linting and type annotations

* make: rm unnecessary SwiGLU

* docs: include modernbert to docs

Author: monatis

docs/modules/index.rst

@@ -6,6 +6,7 @@ API Documentation
 
    attention
    functional
+   models
 
 Attention Modules
 ---------------
@@ -19,6 +20,14 @@ Functional Interfaces
 ------------------
 
 .. automodule:: jax_layers.functional
+   :members:
+   :undoc-members:
+   :show-inheritance: 
+
+Models
+------------------
+
+.. automodule:: jax_layers.models
    :members:
    :undoc-members:
    :show-inheritance: 

jax_layers/models/__init__.py

@@ -0,0 +1,17 @@
+from jax_layers.models.modernbert import (
+    ModernBertAttention,
+    ModernBertEmbeddings,
+    ModernBERTEncoder,
+    ModernBERTForMaskedLM,
+    ModernBertLayer,
+    ModernBertMLP,
+)
+
+__all__ = [
+    "ModernBERTEncoder",
+    "ModernBERTForMaskedLM",
+    "ModernBertAttention",
+    "ModernBertEmbeddings",
+    "ModernBertLayer",
+    "ModernBertMLP",
+]