-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Niru Maheswaranathan
committed
Sep 11, 2015
1 parent
85d16f8
commit 4a0f0cc
Showing
3 changed files
with
114 additions
and
6 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,70 @@ | ||
""" | ||
Tensor | ||
------ | ||
tensor unfolding | ||
""" | ||
|
||
import numpy as np | ||
|
||
__all__ = ['Tensor', 'UnfoldedTensor'] | ||
|
||
|
||
class Tensor(np.ndarray): | ||
|
||
def __new__(cls, arr): | ||
obj = np.asarray(arr).view(cls) | ||
return obj | ||
|
||
def unfold(self, ax): | ||
assert ax in range(self.ndim), "ax less than ndim" | ||
orig_shape = self.shape | ||
rolled_axes = pullax(list(range(self.ndim)), ax) | ||
unfolded = self.transpose(rolled_axes).reshape(orig_shape[ax], -1) | ||
return UnfoldedTensor(unfolded, ax, orig_shape) | ||
|
||
@property | ||
def norm(self): | ||
return np.linalg.norm(self.ravel(), ord=2) | ||
|
||
|
||
class UnfoldedTensor(np.ndarray): | ||
|
||
def __new__(cls, arr, axis, shape): | ||
obj = np.asarray(arr).view(cls) | ||
obj.orig_shape = shape | ||
obj.axis = axis | ||
return obj | ||
|
||
@property | ||
def norm(self): | ||
return np.linalg.norm(self, ord='fro') | ||
|
||
@property | ||
def nucnorm(self): | ||
return np.sum(np.linalg.svd(self, compute_uv=False)) | ||
|
||
def svt(self, threshold): | ||
u, s, v = np.linalg.svd(self, full_matrices=False) | ||
sthr = np.diag(np.maximum(s - threshold, 0)) | ||
return UnfoldedTensor(u.dot(sthr).dot(v), self.axis, self.orig_shape) | ||
|
||
def __array_finalize__(self, obj): | ||
|
||
if obj is None: | ||
print('unfolded tensor None') | ||
return | ||
|
||
self.orig_shape = getattr(obj, 'orig_shape', None) | ||
self.axis = getattr(obj, 'axis', None) | ||
|
||
def fold(self): | ||
rolled_axes = pullax(list(range(len(self.orig_shape))), self.axis) | ||
folded = self.reshape(tuple(self.orig_shape[i] for i in rolled_axes)) | ||
return Tensor(folded.transpose(np.argsort(rolled_axes))) | ||
|
||
|
||
def pullax(values, idx): | ||
values.insert(0, values.pop(idx)) | ||
return tuple(values) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
""" | ||
Tests for the tensor module | ||
""" | ||
|
||
from jetpack.tensor import Tensor | ||
import numpy as np | ||
|
||
|
||
def test_unfolding(): | ||
|
||
t = Tensor(np.random.randn(3, 4, 5)) | ||
|
||
t0 = t.unfold(0) | ||
t1 = t.unfold(1) | ||
t2 = t.unfold(2) | ||
|
||
assert t0.shape == (3, 20) | ||
assert np.allclose(t, t0.fold()) | ||
|
||
assert t1.shape == (4, 15) | ||
assert np.allclose(t, t1.fold()) | ||
|
||
assert t2.shape == (5, 12) | ||
assert np.allclose(t, t2.fold()) | ||
|
||
|
||
def test_norms(): | ||
|
||
# test Frobenius norm | ||
t = Tensor(np.arange(12).reshape(2, 3, 2)) | ||
assert np.allclose(np.linalg.norm(np.arange(12)), t.norm) | ||
assert np.allclose(np.linalg.norm(np.arange(12)), t.unfold(2).norm) | ||
|
||
# test nuclear norm | ||
t = Tensor(np.eye(5, 10).reshape(5, 2, 5)) | ||
assert np.allclose(t.unfold(0).nucnorm, 5.) |