prefab 1.5.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

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/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.dist-info → prefab-1.6.0.dist-info}/METADATA

@@ -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.dist-info → prefab-1.6.0.dist-info}/RECORD

@@ -1,17 +1,17 @@
-prefab/__init__.py,sha256=shfnU_-iDY4MavJaXVb-Xzr9Y7w2_9zhShi4INgEPg4,425
+prefab/__init__.py,sha256=WmddNmU7gmyH4cqkT6BW4iQ3RDG_l5SpKxg4poS6wsU,425
 prefab/__main__.py,sha256=uL6AdCeimPbXiWv0gnq9TDeZhAIrxGwJsEKNYbg9MZg,3454
 prefab/compare.py,sha256=aX7nr9tznSebYeeztvqIPz57npnJ4-iUeKEedrZdksE,3676
 prefab/device.py,sha256=1O6vTOq4wQRGVYvFWLH0uj1XhhYCfnDnIapDEYnBKHw,47996
 prefab/geometry.py,sha256=-nTaGjdw3KN1SVoyvqdcrE2GJP7OqPF6ivUhrO78rUk,11244
-prefab/predict.py,sha256=xEpzyS8uRMZFohBlTGmL90LmVsIn_41ZN7RRhwsw7Ww,18413
+prefab/predict.py,sha256=gxxHCcaMv2TE5O639mQzO3IzFRsLGjY-FQowVp5vLkI,20659
 prefab/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 prefab/read.py,sha256=5BmvFemobA72urxs4j2VZRVvanZZGu1mDB1Uh-epyvI,8635
 prefab/shapes.py,sha256=mRGwsPS-A9XsW3jgvUuMCEeNv9BLXsEnJkytlIHUAKE,28802
 prefab/models/__init__.py,sha256=rRrjcOcHPcob98Coksc0tbvYcXbm6SoLEb-Md233Jvo,1391
 prefab/models/base.py,sha256=t4VNMsOztPedj3kN5fZ1-4tk0SRHWrMuqnIVHztsCs4,1514
 prefab/models/evaluation.py,sha256=2_Klui6tY8xPvOSVD8VpZCVAnT1RX15FONqWG-_x-J8,484
-prefab-1.5.0.dist-info/METADATA,sha256=f8778_o8nPf4DIewUWqsK_2VfImxJish5I390Y9mtww,33754
-prefab-1.5.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-prefab-1.5.0.dist-info/entry_points.txt,sha256=h1_A9O9F3NAIoKXD1RPb3Eo-WCSiHhMB_AnagBi6XTQ,48
-prefab-1.5.0.dist-info/licenses/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-prefab-1.5.0.dist-info/RECORD,,
+prefab-1.6.0.dist-info/METADATA,sha256=CvuP1Khblh_LA-vx5NLgTH5A_zaewTkklkaPEsIY4hI,33754
+prefab-1.6.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+prefab-1.6.0.dist-info/entry_points.txt,sha256=h1_A9O9F3NAIoKXD1RPb3Eo-WCSiHhMB_AnagBi6XTQ,48
+prefab-1.6.0.dist-info/licenses/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
+prefab-1.6.0.dist-info/RECORD,,