Repository: JaxGaussianProcesses/JaxLinOp (language: python)

Issue:
bug: `static_field`s  under `dataclass` inheritance are treated a dynamic.
This behaviour is explained at https://github.com/cgarciae/simple-pytree/issues/8#issue-1630486498

Even though e.g., the `shape` and `dtype` are marked as static on the `LinearOperator` base class, flattening, e.g., a `DenseLinearOperator`, contains these fields.

```python
from jaxlinop import DenseLinearOperator
import jax.numpy as jnp
import jax.tree_util as jtu

jtu.tree_flatten(DenseLinearOperator(jnp.array([[1.0, 2.0], [3.0, 4.0]])))
```
```
([Array([[1., 2.],
         [3., 4.]], dtype=float32),
  2,
  2,
  dtype('float32')],
 PyTreeDef(CustomNode(DenseLinearOperator[(('matrix', 'shape', 'dtype'), {'_pytree__initialized': True})], [*, (*, *), *])))
```

Hopefully this will be resolved in `simple-pytree`, if not we need a work around. To complete this issue we also need to ensure we test for this behaviour.



Discussion / hints:

Fix patch:
diff --git a/setup.py b/setup.py
index 684aa81..4c95eb1 100644
--- a/setup.py
+++ b/setup.py
@@ -37,7 +37,7 @@ def get_versions():
     "jax>=0.4.1",
     "jaxlib>=0.4.1",
     "jaxtyping",
-    "simple-pytree",
+    "simple-pytree>=0.1.6",
 ]
 
 EXTRA_REQUIRE = {


Test patch:
diff --git a/tests/test_linear_operator.py b/tests/test_linear_operator.py
index cbe2082..06a8cf9 100644
--- a/tests/test_linear_operator.py
+++ b/tests/test_linear_operator.py
@@ -14,47 +14,109 @@
 # ==============================================================================
 
 import pytest
+import jax.tree_util as jtu
 import jax.numpy as jnp
 from jaxlinop.linear_operator import LinearOperator
+from dataclasses import dataclass
+from simple_pytree import static_field
 
 
 def test_covariance_operator() -> None:
     with pytest.raises(TypeError):
-        LinearOperator(shape=(1, 1), dtype=jnp.float32)
+        LinearOperator()
 
+@pytest.mark.parametrize("is_dataclass", [True, False])
+@pytest.mark.parametrize("shape", [(1, 1), (2, 3), (4, 5, 6)])
+@pytest.mark.parametrize("dtype", [jnp.float32, jnp.float64])
+def test_instantiate_no_attributes(is_dataclass, shape, dtype) -> None:
+    """Test if the abstract operator can be instantiated, given the abstract methods."""
 
-class DummyLinearOperator(LinearOperator):
-    def diagonal(self, *args, **kwargs):
-        pass
+    class DummyLinearOperator(LinearOperator):
+        def diagonal(self, *args, **kwargs):
+            pass
 
-    def shape(self, *args, **kwargs):
-        pass
+        def shape(self, *args, **kwargs):
+            pass
 
-    def dtype(self, *args, **kwargs):
-        pass
+        def dtype(self, *args, **kwargs):
+            pass
 
-    def __mul__(self, *args, **kwargs):
-        """Multiply linear operator by scalar."""
+        def __mul__(self, *args, **kwargs):
+            """Multiply linear operator by scalar."""
 
-    def _add_diagonal(self, *args, **kwargs):
-        pass
+        def _add_diagonal(self, *args, **kwargs):
+            pass
 
-    def __matmul__(self, *args, **kwargs):
-        """Matrix multiplication."""
+        def __matmul__(self, *args, **kwargs):
+            """Matrix multiplication."""
 
-    def to_dense(self, *args, **kwargs):
-        pass
+        def to_dense(self, *args, **kwargs):
+            pass
 
-    @classmethod
-    def from_dense(self, *args, **kwargs):
-        pass
+        @classmethod
+        def from_dense(cls, *args, **kwargs):
+            pass
 
+    if is_dataclass:
+        dataclass(DummyLinearOperator)
 
-def test_can_instantiate() -> None:
-    """Test if the covariance operator can be instantiated."""
-    res = DummyLinearOperator(shape=(1, 1), dtype=jnp.float32)
+    linop = DummyLinearOperator(shape=shape, dtype=dtype)
+    assert isinstance(linop, DummyLinearOperator)
+    assert isinstance(linop, LinearOperator)
+    assert linop.shape == shape
+    assert linop.dtype == dtype
+    assert jtu.tree_leaves(linop) == [] # shape and dtype are static!
 
-    assert isinstance(res, DummyLinearOperator)
-    assert isinstance(res, LinearOperator)
-    assert res.shape == (1, 1)
-    assert res.dtype == jnp.float32
+
+@pytest.mark.parametrize("is_dataclass", [True, False])
+@pytest.mark.parametrize("shape", [(1, 1), (2, 3), (4, 5, 6)])
+@pytest.mark.parametrize("dtype", [jnp.float32, jnp.float64])
+def test_instantiate_with_attributes(is_dataclass, shape, dtype) -> None:
+    """Test if the covariance operator can be instantiated with attribute annotations."""
+
+    class DummyLinearOperator(LinearOperator):
+        a: int
+        b: int = static_field() # Lets have a static attribute here.
+        c: int
+
+        def __init__(self, shape, dtype, a=1, b=2, c=3):
+            self.a = a
+            self.b = b
+            self.c = c
+            self.shape = shape
+            self.dtype = dtype
+
+        def diagonal(self, *args, **kwargs):
+            pass
+
+        def shape(self, *args, **kwargs):
+            pass
+
+        def dtype(self, *args, **kwargs):
+            pass
+
+        def __mul__(self, *args, **kwargs):
+            """Multiply linear operator by scalar."""
+
+        def _add_diagonal(self, *args, **kwargs):
+            pass
+
+        def __matmul__(self, *args, **kwargs):
+            """Matrix multiplication."""
+
+        def to_dense(self, *args, **kwargs):
+            pass
+
+        @classmethod
+        def from_dense(cls, *args, **kwargs):
+            pass
+
+    if is_dataclass:
+        dataclass(DummyLinearOperator)
+
+    linop = DummyLinearOperator(shape=shape, dtype=dtype)
+    assert isinstance(linop, DummyLinearOperator)
+    assert isinstance(linop, LinearOperator)
+    assert linop.shape == shape
+    assert linop.dtype == dtype
+    assert jtu.tree_leaves(linop) == [1, 3] # b, shape, dtype are static!
\ No newline at end of file