From 7a56e65636bc64b0c1fdda900260c2edf3684ed7 Mon Sep 17 00:00:00 2001
From: BirkhoffG <26811230+BirkhoffG@users.noreply.github.com>
Date: Tue, 28 May 2024 15:50:42 -0700
Subject: [PATCH] Use jnp.clip for applying constraints

---
 nbs/data_utils/transform.ipynb | 2 +-
 relax/data_utils/transforms.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/nbs/data_utils/transform.ipynb b/nbs/data_utils/transform.ipynb
index 46ffd0f..783b3bb 100644
--- 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
index 04971ca..d12bad3 100644
--- 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):