def test_broadcasts_to_multidimensional_arrays(self): tensor = dm_env_rpc_pb2.Tensor() tensor.int32s.array[:] = [4] tensor.shape[:] = [2, 2] unpacked = tensor_utils.unpack_tensor(tensor) expected = np.array([[4, 4], [4, 4]], dtype=np.int32) np.testing.assert_array_equal(expected, unpacked)
def test_integer_broadcasts_1_element_to_all_elements(self): tensor = dm_env_rpc_pb2.Tensor() tensor.int32s.array[:] = [1] tensor.shape[:] = [4] unpacked = tensor_utils.unpack_tensor(tensor) expected = np.array([1, 1, 1, 1], dtype=np.int32) np.testing.assert_array_equal(expected, unpacked)
def test_unpack_multidimensional_arrays(self): tensor = dm_env_rpc_pb2.Tensor() tensor.floats.array[:] = [1, 2, 3, 4, 5, 6, 7, 8] tensor.shape[:] = [2, 4] round_trip = tensor_utils.unpack_tensor(tensor) expected = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]) np.testing.assert_array_equal(expected, round_trip)
def test_negative_dimension_in_matrix(self): tensor = dm_env_rpc_pb2.Tensor() tensor.int32s.array[:] = [1, 2, 3, 4, 5, 6] tensor.shape[:] = [2, -1] unpacked = tensor_utils.unpack_tensor(tensor) expected = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32) np.testing.assert_array_equal(expected, unpacked)
def test_negative_dimension_single_element(self): tensor = dm_env_rpc_pb2.Tensor() tensor.int32s.array[:] = [1] tensor.shape[:] = [-1] unpacked = tensor_utils.unpack_tensor(tensor) expected = np.array([1], dtype=np.int32) np.testing.assert_array_equal(expected, unpacked)
def test_unsigned_integer_broadcasts_1_element_to_all_elements(self): tensor = dm_env_rpc_pb2.Tensor() tensor.uint8s.array = b'\x01' tensor.shape[:] = [4] unpacked = tensor_utils.unpack_tensor(tensor) expected = np.array([1, 1, 1, 1], dtype=np.uint8) np.testing.assert_array_equal(expected, unpacked)
def test_string_broadcasts_1_element_to_all_elements(self): tensor = dm_env_rpc_pb2.Tensor() tensor.strings.array[:] = ['foo'] tensor.shape[:] = [4] unpacked = tensor_utils.unpack_tensor(tensor) expected = np.array(['foo', 'foo', 'foo', 'foo'], dtype=np.str_) np.testing.assert_array_equal(expected, unpacked)
def pack_tensor(value, dtype=None, try_compress=False) -> dm_env_rpc_pb2.Tensor: """Encodes the given value as a tensor. Args: value: A scalar (float, int, string, etc.), protobuf message, NumPy array, or nested lists. dtype: The type to pack the data to. If set to None, will attempt to detect the correct type automatically. Either a dm_env_rpc DataType enum or NumPy type is acceptable. try_compress: A bool, whether to try and encode the tensor in less space or not. This will increase the computational cost of the packing, but may reduce the on-the-wire size of the tensor. There are no guarantees that any compression will actually happen. Raises: ValueError: If `value` is a jagged array, not a primitive type, nested iterable of primitive types or protobuf messages, or all elements can't be cast to the same type or the requested type. Returns: A dm_env_rpc Tensor proto containing the data. """ packed = dm_env_rpc_pb2.Tensor() value = np.asarray(value) if value.dtype == np.object: # Because Numpy doesn't truly support variable-length string arrays, users # tend to use arrays of Numpy objects instead. Iff a user provides an array # of objects and a string dtype argument, automatically convert the value to # an array of strings. if np.issubdtype(_DM_ENV_RPC_DTYPE_TO_NUMPY_DTYPE.get(dtype, dtype), np.str_): for item in value.flat: if not isinstance(item, str): raise TypeError( f'Requested string dtype but not all elements are ' 'Python string types. At least one element was ' f'{type(item)}.') value = np.array(value, dtype=np.str_) elif dtype is not None: value = value.astype(dtype=_DM_ENV_RPC_DTYPE_TO_NUMPY_DTYPE.get( dtype, dtype), copy=False, casting='same_kind') packed.shape[:] = value.shape packer = _TYPE_TO_PACKER[value.dtype.type] if (try_compress and np.all(value == next(value.flat))): # All elements are the same. Pack in to a single value. packer.pack(packed, value.ravel()[0:1]) else: packer.pack(packed, value) return packed
def pack_tensor(value, dtype=None, try_compress=False): """Encodes the given value as a tensor. Args: value: A scalar (float, int, string, etc.), protobuf message, NumPy array, or nested lists. dtype: The type to pack the data to. If set to None, will attempt to detect the correct type automatically. Either a dm_env_rpc DataType enum or NumPy type is acceptable. try_compress: A bool, whether to try and encode the tensor in less space or not. This will increase the computational cost of the packing, but may reduce the on-the-wire size of the tensor. There are no guarantees that any compression will actually happen. Raises: ValueError: If `value` is a jagged array, not a primitive type, nested iterable of primitive types or protobuf messages, or all elements can't be cast to the same type or the requested type. Returns: A dm_env_rpc Tensor proto containing the data. """ packed = dm_env_rpc_pb2.Tensor() value = np.asarray(value) # For efficiency, only check that the first element is a protobuf message. if value.dtype == np.object and value.size > 0 and not isinstance( value.item(0), message.Message): raise ValueError( 'Could not convert value to a tensor of primitive types: ' f'{value}. Are the iterables jagged? Are the data types ' 'not primitive scalar types like strings, floats, or ' 'integers? Or are the elements not protobuf messages?') if dtype is not None: value = value.astype(dtype=_DM_ENV_RPC_DTYPE_TO_NUMPY_DTYPE.get( dtype, dtype), copy=False, casting='same_kind') packed.shape[:] = value.shape packer = _TYPE_TO_PACKER[value.dtype.type] if (try_compress and np.all(value == next(value.flat))): # All elements are the same. Pack in to a single value. packer.pack(packed, value.ravel()[0:1]) else: packer.pack(packed, value) return packed
def test_can_unpack(self): packer = tensor_utils.get_packer(np.int32) tensor = dm_env_rpc_pb2.Tensor() tensor.int32s.array[:] = [1, 2, 3] np.testing.assert_array_equal([1, 2, 3], packer.unpack(tensor))
def test_scalar_with_too_many_elements_raises_error(self): tensor = dm_env_rpc_pb2.Tensor() tensor.int32s.array[:] = [1, 2, 3] with self.assertRaisesRegex(ValueError, '3 element'): tensor_utils.unpack_tensor(tensor)
def test_can_pack(self): packer = tensor_utils.get_packer(np.int32) tensor = dm_env_rpc_pb2.Tensor() packer.pack(tensor, np.asarray([1, 2, 3])) self.assertEqual([1, 2, 3], tensor.int32s.array)
def test_setting_tensor_data_with_wrong_type(self): tensor = dm_env_rpc_pb2.Tensor() with self.assertRaises(TypeError): tensor.floats.array[:] = ['hello!']
def test_two_negative_dimensions_in_matrix(self): tensor = dm_env_rpc_pb2.Tensor() tensor.int32s.array[:] = [1, 2, 3, 4, 5, 6] tensor.shape[:] = [-1, -2] with self.assertRaisesRegex(ValueError, 'one unknown dimension'): tensor_utils.unpack_tensor(tensor)
def test_which_is_set(self): tensor = dm_env_rpc_pb2.Tensor() tensor.floats.array[:] = [1, 2] self.assertEqual('floats', tensor.WhichOneof('payload'))
def test_too_many_elements(self): tensor = dm_env_rpc_pb2.Tensor() tensor.floats.array[:] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] tensor.shape[:] = [2, 4] with self.assertRaisesRegex(ValueError, 'cannot reshape array'): tensor_utils.unpack_tensor(tensor)
def test_setting_tensor_data(self): tensor = dm_env_rpc_pb2.Tensor() tensor.floats.array[:] = [1, 2]