froog 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

froog/__init__.py

@@ -1,3 +1,36 @@
 import froog.optim
 import froog.tensor
-import froog.utils
+import froog.utils
+
+# Import GPU packages
+import froog.gpu.cl.cl_utils
+
+# Try to import Metal utils if available
+try:
+    import froog.gpu.metal.metal_utils
+except ImportError:
+    pass
+
+# Import device management functions
+from froog.gpu import (
+    Device, OpenCLDevice, get_device, set_device,
+    upload_tensor, download_tensor, is_buffer,
+    allocate_buffer, synchronize, get_available_devices
+)
+
+# Try to import Metal device if available
+try:
+    from froog.gpu.metal import MetalDevice
+    __all__ = [
+        'Device', 'OpenCLDevice', 'MetalDevice', 
+        'get_device', 'set_device', 'upload_tensor', 
+        'download_tensor', 'is_buffer', 
+        'allocate_buffer', 'synchronize', 'get_available_devices'
+    ]
+except ImportError:
+    __all__ = [
+        'Device', 'OpenCLDevice', 
+        'get_device', 'set_device', 'upload_tensor', 
+        'download_tensor', 'is_buffer', 
+        'allocate_buffer', 'synchronize', 'get_available_devices'
+    ]

froog/{gradcheck.py → gradient.py}

@@ -7,16 +7,15 @@
 # |___|    |___|  |_||_______||_______||_______|
 
 import numpy as np
+from typing import Callable, Union, Any, Tuple
 from froog.tensor import Tensor
 from froog.utils import mask_like
 
-def jacobian(model, input):
+def jacobian(model: Callable[[Tensor], Tensor], input: Tensor) -> np.ndarray:
   output = model(input)
-
   ji = input.data.reshape(-1).shape[-1]  # jacobian of input
   jo = output.data.reshape(-1).shape[-1] # jacobian of output
   J = np.zeros((jo, ji), dtype=np.float32)
-
   for o in range(jo):
     o_scalar = Tensor(mask_like(output.data, o, 1.)).mul(output).sum()
     o_scalar.backward()
@@ -24,7 +23,7 @@ def jacobian(model, input):
       J[o,i] = grad
   return J
 
-def numerical_jacobian(model, input, eps = 1e-6):
+def numerical_jacobian(model: Callable[[Tensor], Tensor], input: Tensor, eps: float = 1e-6) -> np.ndarray:
 #     """
 #     https://timvieira.github.io/blog/post/2017/04/21/how-to-test-gradient-implementations/
 #     Computes :
@@ -37,24 +36,18 @@ def numerical_jacobian(model, input, eps = 1e-6):
 #         NJ    : an approx. of the Jacobian
 #     """
   output = model(input)
-
   ji = input.data.reshape(-1).shape[-1]
   jo = output.data.reshape(-1).shape[-1]
   NJ = np.zeros((jo, ji), dtype=np.float32)
-
   for i in range(ji):
     eps_perturb = mask_like(input.data, i, mask_value = eps)
-
     output_perturb_add = model(Tensor(input.data + eps_perturb)).data.reshape(-1)
     output_perturb_sub = model(Tensor(input.data - eps_perturb)).data.reshape(-1)
-
     grad_approx = ((output_perturb_add) - (output_perturb_sub)) / (2*eps) # CDM: (f(x + h) - f(x - h)) / (2 * h)
-
     NJ[:,i] = grad_approx
-
   return NJ
 
-def gradcheck(model, input, eps = 1e-06, atol = 1e-5, rtol = 0.001):
+def gradcheck(model: Callable[[Tensor], Tensor], input: Tensor, eps: float = 1e-06, atol: float = 1e-5, rtol: float = 0.001) -> bool:
   """
   Checks whether computed gradient is close to numerical approximation of the Jacobian
   Params: