prefab 1.5.0__tar.gz → 1.6.0__tar.gz

{prefab-1.5.0 → prefab-1.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefab
-Version: 1.5.0
+Version: 1.6.0
 Summary: Artificial nanofabrication of integrated photonic circuits using deep learning
 Project-URL: Homepage, https://prefabphotonics.com
 Project-URL: Repository, https://github.com/PreFab-Photonics/PreFab

{prefab-1.5.0 → prefab-1.6.0}/prefab/__init__.py

@@ -5,7 +5,7 @@ Usage:
     import prefab as pf
 """
 
-__version__ = "1.5.0"
+__version__ = "1.6.0"
 
 from . import compare, geometry, predict, read, shapes
 from .device import BufferSpec, Device

{prefab-1.5.0 → prefab-1.6.0}/prefab/predict.py

@@ -9,6 +9,7 @@ differentiation.
 """
 
 import base64
+import gzip
 import io
 import json
 import os
@@ -384,9 +385,13 @@ def _compute_vjp(
 ) -> npt.NDArray[Any]:
     """Compute J.T @ upstream_gradient via the server-side VJP endpoint."""
     headers = _prepare_headers()
+    upstream_arr = np.squeeze(upstream_gradient).astype(np.float32)
     vjp_data = {
         "device_array": _encode_array(np.squeeze(device_array)),
-        "upstream_gradient": _encode_array(np.squeeze(upstream_gradient)),
+        "upstream_gradient": base64.b64encode(
+            gzip.compress(upstream_arr.tobytes(), compresslevel=1)
+        ).decode("utf-8"),
+        "upstream_gradient_shape": list(upstream_arr.shape),
         "model": model.to_json(),
         "model_type": "p",
     }
@@ -473,6 +478,8 @@ def _predict_array_diff_vjp(
         )
         # Clean up cache
         _diff_cache.pop(cache_key, None)
+        # Ensure gradient shape matches input shape
+        vjp_result = vjp_result.reshape(cached_device_array.shape)
         # Return gradient for device_array, None for model (not differentiable)
         return (vjp_result, None)
 
@@ -487,6 +494,69 @@ predict_array_with_grad = predict_array_diff
 """Alias for predict_array_diff. Deprecated, use predict_array_diff directly."""
 
 
+def differentiable(model: Model):
+    """
+    Create a model-bound differentiable predictor for clean autograd integration.
+
+    Returns a function that takes only `device_array` as input, enabling seamless
+    composition with other differentiable functions. The VJP returns a single
+    gradient array (not a tuple), making it compatible with standard autograd workflows.
+
+    Parameters
+    ----------
+    model : Model
+        The model to use for prediction.
+
+    Returns
+    -------
+    callable
+        A differentiable prediction function that takes `device_array` and returns
+        the predicted fabrication outcome.
+
+    Examples
+    --------
+    >>> predictor = pf.predict.differentiable(model)
+    >>> def loss_fn(x):
+    ...     pred = predictor(x)
+    ...     return np.mean((pred - target) ** 2)
+    >>> gradient = grad(loss_fn)(device_array)  # Returns array, not tuple
+    """
+
+    @primitive
+    def predict(device_array: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        prediction = predict_array(
+            device_array=device_array,
+            model=model,
+            model_type="p",
+            binarize=False,
+        )
+        _diff_cache[id(prediction)] = (device_array.copy(), model)
+        return prediction
+
+    def predict_vjp(
+        ans: npt.NDArray[Any], device_array: npt.NDArray[Any]
+    ) -> Any:
+        cache_key = id(ans)
+        cached_device_array, cached_model = _diff_cache.get(
+            cache_key, (device_array, model)
+        )
+
+        def vjp(g: npt.NDArray[Any]) -> npt.NDArray[Any]:
+            vjp_result = _compute_vjp(
+                device_array=cached_device_array,
+                upstream_gradient=g,
+                model=cached_model,
+            )
+            _diff_cache.pop(cache_key, None)
+            # Ensure gradient shape matches input shape
+            return vjp_result.reshape(device_array.shape)
+
+        return vjp
+
+    defvjp(predict, predict_vjp)
+    return predict
+
+
 def _encode_array(array: npt.NDArray[Any]) -> str:
     """Encode an ndarray as a base64 encoded image for transmission."""
     image = Image.fromarray(np.uint8(array * 255))
@@ -500,7 +570,7 @@ def _decode_array(encoded_png: str) -> npt.NDArray[Any]:
     """Decode a base64 encoded image and return an ndarray."""
     binary_data = base64.b64decode(encoded_png)
     image = Image.open(io.BytesIO(binary_data))
-    return np.array(image) / 255  # type: ignore[no-any-return]
+    return np.array(image) / 255
 
 
 def _prepare_headers() -> dict[str, str]:

{prefab-1.5.0 → prefab-1.6.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "prefab"
-version = "1.5.0"
+version = "1.6.0"
 description = "Artificial nanofabrication of integrated photonic circuits using deep learning"
 authors = [{ name = "PreFab Photonics Inc.", email = "hi@prefabphotonics.com" }]
 keywords = [