[geom] Implement __value_attrs__

tum-pbs · Mar 27, 2024 · f77b24c · f77b24c
1 parent 9520481
commit f77b24c
Show file tree

Hide file tree

Showing 8 changed files with 28 additions and 1 deletion.
diff --git a/PhiML b/PhiML
diff --git a/phi/geom/_box.py b/phi/geom/_box.py
@@ -319,6 +319,9 @@ def __hash__(self):
     def __variable_attrs__(self):
         return '_lower', '_upper'
 
+    def __value_attrs__(self):
+        return ()
+
     @property
     def shape(self):
         if self._lower is None or self._upper is None:

diff --git a/phi/geom/_geom.py b/phi/geom/_geom.py
@@ -549,6 +549,12 @@ def normal(self) -> Tensor:
     def __repr__(self):
         return f"~{self.geometry}"
 
+    def __variable_attrs__(self):
+        return self.geometry.__variable_attrs__
+
+    def __value_attrs__(self):
+        return self.geometry.__value_attrs__
+
 
 def invert(geometry: Geometry):
     """

diff --git a/phi/geom/_geom_ops.py b/phi/geom/_geom_ops.py
@@ -36,6 +36,12 @@ def __init__(self, geometries: Union[Tensor, Geometry]):
     def geometries(self):
         return self._geometries
 
+    def __variable_attrs__(self):
+        return '_geometries',
+
+    def __value_attrs__(self):
+        return '_geometries',
+
     @property
     def object_dims(self):
         return object_dims(self._geometries)

diff --git a/phi/geom/_graph.py b/phi/geom/_graph.py
@@ -22,6 +22,9 @@ def __init__(self, nodes: Union[Geometry, Tensor], connectivity: Tensor, boundar
     def __variable_attrs__(self):
         return '_nodes', '_deltas', '_distances'
 
+    def __value_attrs__(self):
+        return '_nodes',
+
     @property
     def connectivity(self) -> Tensor:
         return self._connectivity

diff --git a/phi/geom/_grid.py b/phi/geom/_grid.py
@@ -211,6 +211,9 @@ def __repr__(self):
     def __variable_attrs__(self):
         return ()
 
+    def __value_attrs__(self):
+        return ()
+
     def __with_attrs__(self, **attrs):
         if not attrs:
             return self

diff --git a/phi/geom/_heightmap.py b/phi/geom/_heightmap.py
@@ -171,6 +171,9 @@ def __repr__(self):
     def __variable_attrs__(self):
         return '_height', '_bounds', '_max_dist', '_fill_below', '_extrapolation', '_faces'
 
+    def __value_attrs__(self):
+        return ()
+
     def __getitem__(self, item):
         item = slicing_dict(self, item)
         return Heightmap(self._height[item], self._bounds[item], self._max_dist[item], self._fill_below[item], self._extrapolation[item] if self._extrapolation is not None else None, math.slice(self._faces, item))

diff --git a/phi/geom/_sphere.py b/phi/geom/_sphere.py
@@ -104,6 +104,9 @@ def scaled(self, factor: Union[float, Tensor]) -> 'Geometry':
     def __variable_attrs__(self):
         return '_center', '_radius'
 
+    def __value_attrs__(self):
+        return '_center',
+
     def __getitem__(self, item):
         item = slicing_dict(self, item)
         return Sphere(self._center[_keep_vector(item)], self._radius[item])
+1 −1		.github/workflows/unit-tests.yml
+20 −3		README.md
+929 −0		docs/Advantages_Data_Types.ipynb
+142 −128		docs/Data_Types.ipynb
+616 −0		docs/Dimension_Names_Types.ipynb
+0 −397		docs/Networks.ipynb
+408 −0		docs/Performance.ipynb
+417 −359		docs/Shapes.ipynb
+145 −119		docs/Tensors.ipynb
+6 −0		docs/index.md
+1 −1		phiml/VERSION
+71 −4		phiml/_troubleshoot.py
+2 −2		phiml/backend/__init__.py
+141 −37		phiml/backend/_backend.py
+1 −1		phiml/backend/_dtype.py
+66 −57		phiml/backend/_linalg.py
+39 −18		phiml/backend/_minimize.py
+69 −24		phiml/backend/_numpy_backend.py
+63 −33		phiml/backend/jax/_jax_backend.py
+18 −15		phiml/backend/jax/stax_nets.py
+134 −75		phiml/backend/tensorflow/_tf_backend.py
+14 −11		phiml/backend/tensorflow/nets.py
+122 −65		phiml/backend/torch/_torch_backend.py
+22 −20		phiml/backend/torch/nets.py
+12 −7		phiml/math/__init__.py
+159 −86		phiml/math/_functional.py
+45 −17		phiml/math/_magic_ops.py
+199 −16		phiml/math/_nd.py
+459 −193		phiml/math/_ops.py
+100 −56		phiml/math/_optimize.py
+116 −57		phiml/math/_shape.py
+394 −65		phiml/math/_sparse.py
+161 −116		phiml/math/_tensors.py
+550 −84		phiml/math/_trace.py
+70 −13		phiml/math/extrapolation.py
+10 −6		phiml/math/magic.py
+40 −36		phiml/nn.py
+2 −1		tests/commit/backend/test__backend.py
+1 −0		tests/commit/backend/test_jax_backend.py
+1 −0		tests/commit/backend/test_tf_backend.py
+1 −0		tests/commit/backend/test_torch_backend.py
+63 −4		tests/commit/math/test__functional.py
+79 −0		tests/commit/math/test__nd.py
+123 −0		tests/commit/math/test__ops.py
+12 −0		tests/commit/math/test__optimize.py
+25 −1		tests/commit/math/test__shape.py
+26 −0		tests/commit/math/test__sparse.py
+12 −0		tests/commit/math/test__tensors.py
+13 −1		tests/commit/math/test__trace.py