temp

Author: Alexander

examples/Bert.py

@@ -55,12 +55,12 @@ def __init__(
 
     def __call__(
         self,
-        token_ids: Int[Array, "seq_len"],
-        position_ids: Int[Array, "seq_len"],
-        segment_ids: Int[Array, "seq_len"],
+        token_ids: Array,
+        position_ids: Array,
+        segment_ids: Array,
         enable_dropout: bool = False,
         key: jax.random.PRNGKey | None = None,
-    ) -> Float[Array, "seq_len hidden_size"]:
+    ) -> Array:
         tokens = jax.vmap(self.token_embedder)(token_ids)
         segments = jax.vmap(self.segment_embedder)(segment_ids)
         positions = jax.vmap(self.position_embedder)(position_ids)
@@ -124,9 +124,9 @@ def __init__(
 
     def __call__(
         self,
-        token_ids: Int[Array, "seq_len"],
-        position_ids: Int[Array, "seq_len"],
-        segment_ids: Int[Array, "seq_len"],
+        token_ids: Array,
+        position_ids: Array,
+        segment_ids: Array,
         *,
         enable_dropout: bool = False,
         key: jax.random.PRNGKey | None = None,
@@ -190,10 +190,10 @@ def __init__(self, config: Mapping, num_classes: int, key: jax.random.PRNGKey):
 
     def __call__(
         self,
-        inputs: dict[str, Int[Array, "seq_len"]],
+        inputs: dict[str, Array],
         enable_dropout: bool = True,
         key: jax.random.PRNGKey = None,
-    ) -> Float[Array, "num_classes"]:
+    ) -> Array:
         seq_len = inputs["token_ids"].shape[-1]
         position_ids = jnp.arange(seq_len)
 

examples/transformer.py

@@ -93,11 +93,11 @@ def __init__(
 
     def __call__(
         self,
-        inputs: Float[Array, "seq_len hidden_size"],
-        mask: Int[Array, " seq_len"] | None,
+        inputs: Array,
+        mask: Array | None,
         enable_dropout: bool = False,
         key: "jax.random.PRNGKey" = None,
-    ) -> Float[Array, "seq_len hidden_size"]:
+    ) -> Array:
         if mask is not None:
             mask = self.make_self_attention_mask(mask)
         attention_key, dropout_key = (
@@ -132,7 +132,7 @@ def _single_head_attention(q, k, v):
         return result
 
     def make_self_attention_mask(
-        self, mask: Int[Array, " seq_len"]
+        self, mask: Array
     ) -> Float[Array, "num_heads seq_len seq_len"]:
         """Create self-attention mask from sequence-level mask."""
         mask = jnp.multiply(
@@ -176,12 +176,12 @@ def __init__(
 
     def __call__(
         self,
-        inputs: Float[Array, "seq_len hidden_size"],
-        mask: Int[Array, " seq_len"] | None = None,
+        inputs: Array,
+        mask: Array | None = None,
         *,
         enable_dropout: bool = False,
         key: jax.random.PRNGKey | None = None,
-    ) -> Float[Array, "seq_len hidden_size"]:
+    ) -> Array:
         attn_key, ff_key = (None, None) if key is None else jax.random.split(key)
         attention_output = self.attention_block(
             inputs, mask, enable_dropout=enable_dropout, key=attn_key