Make materialize public to the transform package

- materialize is in a protected file _materialize.py, which means that only members of the local transform package can import from it. There is thus no point in make the materialize function itself protected. In fact, this would mean that the materialize function should not be used outside _materialize.py, while in fact, it has to be used in grad.py, jac.py, and stack.py.

Author: ValerianRey

src/torchjd/autojac/_transform/_materialize.py

@@ -4,7 +4,7 @@
 from torch import Tensor
 
 
-def _materialize(
+def materialize(
     optional_tensors: Sequence[Tensor | None], inputs: Sequence[Tensor]
 ) -> tuple[Tensor, ...]:
     """

src/torchjd/autojac/_transform/grad.py

@@ -4,7 +4,7 @@
 from torch import Tensor
 
 from ._differentiate import _Differentiate
-from ._materialize import _materialize
+from ._materialize import materialize
 from .tensor_dict import Gradients
 
 
@@ -47,5 +47,5 @@ def _differentiate(self, grad_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]:
             create_graph=self.create_graph,
             allow_unused=True,
         )
-        grads = _materialize(optional_grads, inputs)
+        grads = materialize(optional_grads, inputs)
         return grads

src/torchjd/autojac/_transform/jac.py

@@ -7,7 +7,7 @@
 from torch import Size, Tensor
 
 from ._differentiate import _Differentiate
-from ._materialize import _materialize
+from ._materialize import materialize
 from .tensor_dict import Jacobians
 
 
@@ -60,7 +60,7 @@ def _get_vjp(grad_outputs: Sequence[Tensor], retain_graph: bool) -> Tensor:
                 create_graph=self.create_graph,
                 allow_unused=True,
             )
-            grads = _materialize(optional_grads, inputs=inputs)
+            grads = materialize(optional_grads, inputs=inputs)
             return torch.concatenate([grad.reshape([-1]) for grad in grads])
 
         # By the Jacobians constraint, this value should be the same for all jac_outputs.

src/torchjd/autojac/_transform/stack.py

@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from ._materialize import _materialize
+from ._materialize import materialize
 from .base import Transform
 from .tensor_dict import _A, Gradients, Jacobians
 
@@ -46,6 +46,6 @@ def _stack_one_key(gradient_dicts: list[Gradients], input: Tensor) -> Tensor:
     """
 
     optional_gradients = [gradients.get(input, None) for gradients in gradient_dicts]
-    gradients = _materialize(optional_gradients, [input] * len(optional_gradients))
+    gradients = materialize(optional_gradients, [input] * len(optional_gradients))
     jacobian = torch.stack(gradients, dim=0)
     return jacobian