[FIX] amplitude and ls property types (#32)

a-ws-m · Oct 22, 2021 · ba6c627 · ba6c627
1 parent 7072967
commit ba6c627
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/tests/test_kernels.py b/tests/test_kernels.py
@@ -58,7 +58,7 @@ def test_reload(tmp_path: Path, kernel_type: Type[AmpAndLengthScaleFn]):
     loaded_weights = loaded_kernel.get_weights()
     assert loaded_weights == example_weights
 
-    loaded_amp = loaded_kernel.amplitude
-    loaded_ls = loaded_kernel.length_scale
+    loaded_amp = loaded_kernel.amplitude.numpy()
+    loaded_ls = loaded_kernel.length_scale.numpy()
     assert loaded_amp == pytest.approx(orig_amplitude)
     assert loaded_ls == pytest.approx(orig_length_scale)
diff --git a/unlocknn/kernel_layers.py b/unlocknn/kernel_layers.py
@@ -110,14 +110,14 @@ def __init__(self, **kwargs):
         )
 
     @property
-    def amplitude(self) -> float:
+    def amplitude(self) -> tf.Tensor:
         """Get the current kernel amplitude."""
-        return tf.nn.softplus(0.1 * self._amplitude_basis).numpy().item()
+        return tf.nn.softplus(0.1 * self._amplitude_basis)
 
     @property
-    def length_scale(self) -> float:
+    def length_scale(self) -> tf.Tensor:
         """Get the current kernel length scale."""
-        return tf.nn.softplus(5.0 * self._length_scale_basis).numpy().item()
+        return tf.nn.softplus(5.0 * self._length_scale_basis)
 
 
 class RBFKernelFn(AmpAndLengthScaleFn):