fix: rename equal_shape

Signed-off-by: Aaron Pham <[email protected]>
bentoml · Feb 17, 2023 · 300959a · 300959a
1 parent 8ad79c9
commit 300959a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/src/bentoml/_internal/frameworks/flax.py b/src/bentoml/_internal/frameworks/flax.py
@@ -62,7 +62,7 @@
 
 @attr.define
 class FlaxOptions(ModelOptions):
-    """Options for the Keras model."""
+    """Options for the Flax model."""
 
     partial_kwargs: t.Dict[str, t.Any] = attr.field(factory=dict)
     device: str = attr.field(default="cpu")

diff --git a/tests/integration/frameworks/models/flax.py b/tests/integration/frameworks/models/flax.py
@@ -61,7 +61,9 @@ def init_mlp_state():
     return params
 
 
-def is_close(model: nn.Module, state_dict: dict[str, t.Any], arr: jnp.ndarray):
+def assert_equal_shape(
+    model: nn.Module, state_dict: dict[str, t.Any], arr: jnp.ndarray
+):
     def check(out: jnp.ndarray) -> bool:
         logit = model.apply({"params": state_dict["params"]}, arr)
         chex.assert_equal_shape([logit, out])
@@ -84,7 +86,9 @@ def check(out: jnp.ndarray) -> bool:
                 "__call__": [
                     Input(
                         input_args=[ones_array],
-                        expected=is_close(MLP(), init_mlp_state(), ones_array),
+                        expected=assert_equal_shape(
+                            MLP(), init_mlp_state(), ones_array
+                        ),
                     )
                 ]
             },