Use jnp.clip for applying constraints

BirkhoffG · May 28, 2024 · 7a56e65 · 7a56e65
1 parent c722a81
commit 7a56e65
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/nbs/data_utils/transform.ipynb b/nbs/data_utils/transform.ipynb
@@ -158,7 +158,7 @@
     "        super().__init__(\"minmax\", MinMaxScaler())\n",
     "\n",
     "    def apply_constraints(self, xs, cfs, **kwargs):\n",
-    "        return np.clip(cfs, 0., 1.)"
+    "        return jnp.clip(cfs, 0., 1.)"
    ]
   },
   {

diff --git a/relax/data_utils/transforms.py b/relax/data_utils/transforms.py
@@ -92,7 +92,7 @@ def __init__(self):
         super().__init__("minmax", MinMaxScaler())
 
     def apply_constraints(self, xs, cfs, **kwargs):
-        return np.clip(cfs, 0., 1.)
+        return jnp.clip(cfs, 0., 1.)
 
 # %% ../../nbs/data_utils/transform.ipynb 9
 class _OneHotTransformation(_DefaultTransformation):