refactor: Update types of all models #222

mehmetozsoy-synnada · 2025-03-05T17:13:21Z

Description

Closes #221.

What is Changed

Include the changes introduced in this PR.

cutoff -> threshold
types of all primitives are updated
all instances of general_tensor_type_constraint is replaced with general_type_constraint
values that are taken from __call__() carried to __init__()
tests are updated accordingly.

Checklist:

Tests that cover the code added.
Corresponding changes documented.
All tests passed.
The code linted and styled (pre-commit run --all-files has passed).

…Left and ShiftRight is made polymorphic, initial_indexer_type_constraint is updated

…into update-types-of-all-models

…are updated, tests are fixed

…into update-types-of-all-models

codecov · 2025-03-05T18:11:04Z

Codecov Report

Attention: Patch coverage is 92.85714% with 7 lines in your changes missing coverage. Please review.

Project coverage is 89.31%. Comparing base (8938a63) to head (0350620).

Files with missing lines	Patch %	Lines
mithril/models/primitives.py	87.80%	5 Missing ⚠️
mithril/cores/python/numpy/ops_grad.py	92.85%	1 Missing ⚠️
mithril/cores/python/torch/ops.py	87.50%	1 Missing ⚠️

@@           Coverage Diff           @@
##             main     #222   +/-   ##
=======================================
  Coverage   89.30%   89.31%           
=======================================
  Files          68       68           
  Lines       16666    16578   -88     
=======================================
- Hits        14884    14807   -77     
+ Misses       1782     1771   -11

Files with missing lines	Coverage Δ
mithril/cores/python/common_primitives.py	`67.18% <100.00%> (+0.52%)`	⬆️
mithril/cores/python/jax/ops.py	`85.51% <100.00%> (-0.09%)`	⬇️
mithril/cores/python/mlx/ops.py	`86.40% <100.00%> (ø)`
mithril/cores/python/numpy/ops.py	`81.40% <100.00%> (ø)`
mithril/framework/constraints.py	`94.04% <100.00%> (+0.40%)`	⬆️
mithril/framework/logical/operators.py	`100.00% <100.00%> (ø)`
mithril/models/train_model.py	`91.88% <ø> (ø)`
mithril/cores/python/numpy/ops_grad.py	`88.12% <92.85%> (ø)`
mithril/cores/python/torch/ops.py	`80.73% <87.50%> (-0.11%)`	⬇️
mithril/models/primitives.py	`97.91% <87.80%> (-0.11%)`	⬇️

... and 2 files with indirect coverage changes

kberat-synnada · 2025-03-07T13:38:06Z

mithril/cores/python/jax/ops.py

-def robust_sqrt(input: jax.Array, cutoff: jax.Array) -> jax.Array:
-    # v_mapped_func= jax.vmap(partial(robust_log_helper, threshold = cutoff))
+def robust_sqrt(input: jax.Array, threshold: jax.Array) -> jax.Array:
+    # v_mapped_func= jax.vmap(partial(robust_log_helper, threshold = threshold))


we can remove unnecessary comments

kberat-synnada · 2025-03-07T13:38:14Z

mithril/cores/python/jax/ops.py

-def robust_log(input: jax.Array, cutoff: jax.Array) -> jax.Array:
-    # v_mapped_func= jax.vmap(partial(robust_log_helper, threshold = cutoff))
+def robust_log(input: jax.Array, threshold: jax.Array) -> jax.Array:
+    # v_mapped_func= jax.vmap(partial(robust_log_helper, threshold = threshold))


we can remove unnecessary comments

kberat-synnada · 2025-03-07T13:44:51Z

mithril/framework/constraints.py

@@ -142,6 +143,14 @@ def is_repr_known(repr: ShapeRepr) -> bool:
    return repr.root is None and all([uni.value is not None for uni in repr.prefix])


+def two_add(input1: Any, input2: Any, input3: Any) -> Any:


we can convert this function to accept *args as input and sum them all for more generic purposes

kberat-synnada · 2025-03-07T13:46:43Z

mithril/models/primitives.py

-            input=BaseKey(shape=[("Var_2", ...)], type=Tensor, value=input),
-            target=BaseKey(shape=[("Var_3", ...)], type=Tensor, value=target),
-            cutoff=BaseKey(shape=[], type=Tensor, value=cutoff),
+            input=BaseKey(


no need to set shapes explicitly if they are different Variadics. It is the default behaviour

kberat-synnada · 2025-03-07T13:48:23Z

mithril/models/primitives.py

@@ -890,6 +901,7 @@ def __init__(
            approximate=BaseKey(value=approximate, type=bool),
            name=name,
            input=input,
+            output=BaseKey(shape=[("Var", ...)], type=Tensor[float]),


no need to set shapes

kberat-synnada · 2025-03-07T13:48:59Z

mithril/models/primitives.py

+            name=name,
+            axis=axis_key,
+            input=input,
+            output=BaseKey(shape=[("Var", ...)], type=Tensor[float]),


no need to set shapes explicitly

kberat-synnada · 2025-03-07T13:49:08Z

mithril/models/primitives.py

+            slope=BaseKey(
+                value=slope, type=int | float | bool | Tensor[int | float | bool]
+            ),
+            output=BaseKey(shape=[("Var", ...)], type=Tensor[float]),


no need to set shapes explicitly

kberat-synnada · 2025-03-07T13:50:56Z

mithril/models/primitives.py

@@ -2109,7 +2164,7 @@ def __init__(
        super().__init__(
            formula_key="randn",
            name=name,
-            output=BaseKey(shape=[("output", ...)], type=Tensor),
+            output=BaseKey(shape=[("output", ...)], type=Tensor[float]),


no need to set shapes explicitly

kberat-synnada · 2025-03-07T14:02:16Z

We should be able to remove wrapping operation in PhysicalModel, this part:

        constant_keys = {  # type: ignore
            key: StaticDataStore.convert_data_to_physical(value, backend)  # type: ignore
            for key, value in model().connections.items()
            if value is not NOT_GIVEN
        } | constant_keys

If we can get rid of this, we can delete convert_data_to_physical too

mehmetozsoy-synnada added 20 commits February 12, 2025 17:27

general_type_constraint code is written

69145c1

general_type_constraint is updated, operators are updated

3966a39

merged with main

ce8df3f

scalar value tests are updated, type inference tests are added, Shift…

417c6a8

…Left and ShiftRight is made polymorphic, initial_indexer_type_constraint is updated

merged with main

43e466f

apply ruff updates

de366e0

minor test bugs is fixed

e4cecd5

edge_type_constraints is removed

107d667

reviews are updated

5f80323

add type inference_tests

d3a1142

merged with upstream

3e634ea

reviews are updated

5c8f4f3

edge referee update type logic is updated

c8f1c1a

types are updated

14dae18

merged with main

12fdc29

types are partially updated

d6b656e

init if primitive models are partially updated

b536375

Merge branch 'main' of https://github.com/mehmetozsoy-synnada/mithril …

ac488f1

…into update-types-of-all-models

all values carried from call to init, cutoff -> threshold, all types …

25f3999

…are updated, tests are fixed

Merge branch 'main' of https://github.com/mehmetozsoy-synnada/mithril …

0350620

…into update-types-of-all-models

github-actions bot added tests models constraints examples labels Mar 5, 2025

github-actions bot assigned norhan-synnada and kberat-synnada Mar 5, 2025

kberat-synnada reviewed Mar 7, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

refactor: Update types of all models #222

refactor: Update types of all models #222

mehmetozsoy-synnada commented Mar 5, 2025

codecov bot commented Mar 5, 2025

kberat-synnada Mar 7, 2025

kberat-synnada Mar 7, 2025

kberat-synnada Mar 7, 2025

kberat-synnada Mar 7, 2025

kberat-synnada Mar 7, 2025

kberat-synnada Mar 7, 2025

kberat-synnada Mar 7, 2025

kberat-synnada Mar 7, 2025

kberat-synnada commented Mar 7, 2025

		@@ -142,6 +143,14 @@ def is_repr_known(repr: ShapeRepr) -> bool:
		return repr.root is None and all([uni.value is not None for uni in repr.prefix])


		def two_add(input1: Any, input2: Any, input3: Any) -> Any:

refactor: Update types of all models #222

Are you sure you want to change the base?

refactor: Update types of all models #222

Conversation

mehmetozsoy-synnada commented Mar 5, 2025

Description

What is Changed

Checklist:

codecov bot commented Mar 5, 2025

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kberat-synnada commented Mar 7, 2025