Change tensor_spec_utils.bounds to not broadcast scalar bounds.

PiperOrigin-RevId: 371299720
Change-Id: I166d2a8d42e0544eec0a4e425181d391e3fe7410

Author: tomwardio

dm_env_rpc/v1/tensor_spec_utils.py

@@ -65,12 +65,13 @@ def _get_value(min_max_value, shape, default):
   value = which and getattr(min_max_value, which)
 
   if value is None:
-    min_max = np.broadcast_to(default, shape) if shape else default
+    min_max = default
   elif which in _SCALAR_VALUE_TYPES:
-    min_max = np.broadcast_to(value, shape) if shape else value
+    min_max = value
   else:
     unpacked = tensor_utils.unpack_proto(min_max_value)
-    min_max = tensor_utils.reshape_array(unpacked, shape)
+    min_max = tensor_utils.reshape_array(
+        unpacked, shape) if len(unpacked) > 1 else unpacked[0]
 
   if (shape is not None
       and np.any(np.array(shape) < 0)

dm_env_rpc/v1/tensor_spec_utils_test.py

@@ -62,7 +62,7 @@ def test_min_n_shape(self):
     tensor_spec.shape[:] = minimum.shape
     bounds = tensor_spec_utils.bounds(tensor_spec)
     np.testing.assert_array_equal(minimum, bounds.min)
-    np.testing.assert_array_equal(np.full(minimum.shape, 2**32 - 1), bounds.max)
+    self.assertEqual(2**32 - 1, bounds.max)
 
   def test_max_n_shape(self):
     maximum = np.array([[1, 2], [3, 4]])
@@ -72,7 +72,7 @@ def test_max_n_shape(self):
     tensor_spec.max.uint32s.array[:] = maximum.flatten().data.tolist()
     tensor_spec.shape[:] = maximum.shape
     bounds = tensor_spec_utils.bounds(tensor_spec)
-    np.testing.assert_array_equal(np.full(maximum.shape, 0), bounds.min)
+    self.assertEqual(0, bounds.min)
     np.testing.assert_array_equal(maximum, bounds.max)
 
   def test_invalid_min_shape(self):
@@ -134,34 +134,33 @@ def test_invalid_min_var_shape(self):
         'can only have scalar ranges.'):
       tensor_spec_utils.bounds(tensor_spec)
 
-  def test_min_broadcast(self):
+  def test_min_scalar_doesnt_broadcast(self):
     tensor_spec = dm_env_rpc_pb2.TensorSpec()
     tensor_spec.dtype = dm_env_rpc_pb2.DataType.UINT32
     tensor_spec.min.uint32s.array[:] = [1]
     tensor_spec.shape[:] = (2, 2)
     bounds = tensor_spec_utils.bounds(tensor_spec)
-    np.testing.assert_array_equal(np.full(tensor_spec.shape, 1), bounds.min)
-    np.testing.assert_array_equal(
-        np.full(tensor_spec.shape, 2**32 - 1), bounds.max)
+    self.assertEqual(1, bounds.min)
+    self.assertEqual(2**32 - 1, bounds.max)
 
-  def test_max_broadcast(self):
+  def test_max_scalar_doesnt_broadcast(self):
     tensor_spec = dm_env_rpc_pb2.TensorSpec()
     tensor_spec.dtype = dm_env_rpc_pb2.DataType.UINT32
     tensor_spec.max.uint32s.array[:] = [1]
     tensor_spec.shape[:] = (2, 2)
     bounds = tensor_spec_utils.bounds(tensor_spec)
-    np.testing.assert_array_equal(np.full(tensor_spec.shape, 0), bounds.min)
-    np.testing.assert_array_equal(np.full(tensor_spec.shape, 1), bounds.max)
+    self.assertEqual(0, bounds.min)
+    self.assertEqual(1, bounds.max)
 
-  def test_min_max_broadcast(self):
+  def test_min_max_scalars_dont_broadcast(self):
     tensor_spec = dm_env_rpc_pb2.TensorSpec()
     tensor_spec.dtype = dm_env_rpc_pb2.DataType.UINT32
     tensor_spec.min.uint32s.array[:] = [1]
     tensor_spec.max.uint32s.array[:] = [2]
     tensor_spec.shape[:] = (4,)
     bounds = tensor_spec_utils.bounds(tensor_spec)
-    np.testing.assert_array_equal(np.full(tensor_spec.shape, 1), bounds.min)
-    np.testing.assert_array_equal(np.full(tensor_spec.shape, 2), bounds.max)
+    self.assertEqual(1, bounds.min)
+    self.assertEqual(2, bounds.max)
 
   def test_min_mismatches_type_raises_error(self):
     tensor_spec = dm_env_rpc_pb2.TensorSpec()