def test_nested_map():
func = lambda x: x * 2 if isinstance(x, (
int, float)) else x + ' Not!' if isinstance(x, basestring) else x
assert nested_map(func, 2) == 4
assert nested_map(func, 'God is dead.') == 'God is dead. Not!'
assert nested_map(func, (1, 2, 3)) == (2, 4, 6)
assert nested_map(func, [1, 2, None, {
'a': 3,
'b': 'It works!'
}]) == [2, 4, None, {
'a': 6,
'b': 'It works! Not!'
}]
def to_structseq(self, as_arrays=False):
structs = self.to_struct()
return nested_map(
lambda s: seqstruct_to_structseq(s, as_arrays=as_arrays)
if isinstance(s, list) else s,
structs,
is_container_func=lambda x: isinstance(x, dict))
def test_nested_map_with_container_func():
data = {'a': [1, 2, 3], 'b': [3, 4, 5]}
result = nested_map(lambda x: np.array(x), data, is_container_func=lambda x: isinstance(x, dict))
assert isinstance(result['a'], np.ndarray)
assert np.array_equal(result['a'], [1, 2, 3])
assert isinstance(result['b'], np.ndarray)
assert np.array_equal(result['b'], [3, 4, 5])
def test_nested_map_with_container_func():
data = {'a': [1, 2, 3], 'b': [3, 4, 5]}
result = nested_map(lambda x: np.array(x), data, is_container_func=lambda x: isinstance(x, dict))
assert isinstance(result['a'], np.ndarray)
assert np.array_equal(result['a'], [1, 2, 3])
assert isinstance(result['b'], np.ndarray)
assert np.array_equal(result['b'], [3, 4, 5])
def cleave(self, other = None, share_data=True, to_params = False):
if other is not None:
assert other.size() == self.size()
else:
other = self
split_vars = self._join_split.cleave(other, share_data=share_data)
if to_params:
split_vars = nested_map(lambda v: Parameter(v.data), split_vars)
return split_vars
def forward(self, x, is_cuda=0):
if self.state is None:
self._initialize_state(x)
if is_cuda:
self.state = self.state.cuda()
out, self.state = self.forward_update_module(
x, nested_map(lambda x: x.detach(), self.state))
return out
def torch_str(var, thresh=10):
if isinstance(var, (Variable, )):
if np.prod(var.size()) > thresh:
return '<{} with size {}, min {}, max {}>'.format(
var.__class__.__name__, var.size(),
var.min().data[0],
var.max().data[0])
else:
# return '<{} with values {}>'.format(var.__class__.__name__, var.data.numpy())
return str(var.data.numpy())
elif isinstance(var, (list, tuple, dict)):
return nested_map(torch_str, var)
else:
raise Exception(var.__class__)
def test_nested_map():
func = lambda x: x*2 if isinstance(x, (int, float)) else x+' Not!' if isinstance(x, string_types) else x
assert nested_map(func, 2)==4
assert nested_map(func, 'God is dead.')=='God is dead. Not!'
assert nested_map(func, (1, 2, 3)) == (2, 4, 6)
assert nested_map(func, [1, 2, None, {'a': 3, 'b': 'It works!'}]) == [2, 4, None, {'a': 6, 'b': 'It works! Not!'}]
with pytest.raises(AssertionError):
assert nested_map(lambda a, b: a+b, {'a': 1, 'b': [2, 3]}, {'a': 4, 'XXX': [5, 6]}) == {'a': 5, 'b': [7, 9]}
with pytest.raises(AssertionError):
assert nested_map(lambda a, b: a+b, {'a': 1, 'b': [2, 3]}, {'a': 4, 'b': [5, 6], 'c': [7]}) == {'a': 5, 'b': [7, 9]}
def test_nested_map():
func = lambda x: x*2 if isinstance(x, (int, float)) else x+' Not!' if isinstance(x, string_types) else x
assert nested_map(func, 2)==4
assert nested_map(func, 'God is dead.')=='God is dead. Not!'
assert nested_map(func, (1, 2, 3)) == (2, 4, 6)
assert nested_map(func, [1, 2, None, {'a': 3, 'b': 'It works!'}]) == [2, 4, None, {'a': 6, 'b': 'It works! Not!'}]
with pytest.raises(AssertionError):
assert nested_map(lambda a, b: a+b, {'a': 1, 'b': [2, 3]}, {'a': 4, 'XXX': [5, 6]}) == {'a': 5, 'b': [7, 9]}
with pytest.raises(AssertionError):
assert nested_map(lambda a, b: a+b, {'a': 1, 'b': [2, 3]}, {'a': 4, 'b': [5, 6], 'c': [7]}) == {'a': 5, 'b': [7, 9]}
def to_seqstruct(self):
structs = self.to_struct()
return nested_map(lambda s: seqstruct_to_structseq(s)
if isinstance(s, dict) else s,
structs,
is_container_func=lambda x: isinstance(x, list))
def numpy_struct_to_torch_struct(obj, cast_floats=None):
return nested_map(partial(_numpy_to_torch_var, cast_floats=cast_floats),
obj)
def torch_struct_to_numpy_struct(obj):
return nested_map(_torch_to_numpy, obj)
def clone_em(vars):
return nested_map(lambda x: x.clone()
if isinstance(x, Variable) else x, vars)
