make boxes non-diff-able

hanyas · Aug 3, 2022 · c22c40a · c22c40a
1 parent 830adc4
commit c22c40a
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/tox/objects.py b/tox/objects.py
@@ -2,10 +2,12 @@
 import jax.numpy as jnp
 
 
-class Box(NamedTuple):
- low: jnp.ndarray
- high: jnp.ndarray
- shape: Tuple
+class Box:
+
+ def __init__(self, low: jnp.ndarray, high: jnp.ndarray, shape: Tuple):
+ self.low = low
+ self.high = high
+ self.shape = shape
 
  def clip(self, x: jnp.ndarray) -> jnp.ndarray:
  return jnp.clip(x, self.low, self.high)